big hard cwm

app.py

@@ -26,7 +26,7 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 
 # Load CWM 3-frame model (automatically download pre-trained checkpoint)
-model = model_factory.load_model('vitb_8x8patch_3frames')#.to(device)
+model = model_factory.load_model('vitb_8x8patch_2frames_encoder_mask_token')#.to(device)
 
 model.requires_grad_(False)
 model.eval()
@@ -89,7 +89,7 @@ import os
 # preloaded_images = load_preuploaded_images()
 #
 # print("Preloaded images:", preloaded_images)
-@spaces.GPU(duration=110)
+# @spaces.GPU(duration=110)
 def get_c(x, points):
     x = utils.imagenet_normalize(x)#.to(device)
     with torch.no_grad():
@@ -107,6 +107,7 @@ with gr.Blocks() as demo:
             with gr.Column():
                 # Input image
                 original_image = gr.State(value=None)  # store original image without arrows
+                original_image_high_res = gr.State(value=None)  # store original image without arrows
                 input_image = gr.Image(type="numpy", label="Upload Image")
 
                 # Annotate arrows
@@ -125,12 +126,12 @@ with gr.Blocks() as demo:
                 output_image = gr.Image(type='numpy')
 
         # Store the original image and resize to square size once uploaded
-        def resize_to_square(img, size=448):
+        def resize_to_square(img, size=512):
             print("Resizing image to square")
             img = Image.fromarray(img)
             transform = transforms.Compose([
                 transforms.Resize((size, size)),
-                transforms.CenterCrop(size)
+                # transforms.CenterCrop(size)
             ])
             img = transform(img)  # .transpose(1, 2, 0)
 
@@ -142,13 +143,13 @@ with gr.Blocks() as demo:
             img = np.array(Image.open(img_path))
             # print(f"Image uploaded with shape: {input.shape}")
             resized_img = resize_to_square(img)
-            return resized_img, resized_img, []
+            return resized_img, resized_img, img,  []
 
 
         def store_img(img):
             resized_img = resize_to_square(img)  # Resize the uploaded image to a square
             print(f"Image uploaded with shape: {resized_img.shape}")
-            return resized_img, resized_img, []
+            return resized_img, resized_img, img, []
 
 
         with gr.Row():
@@ -165,9 +166,9 @@ with gr.Blocks() as demo:
                 #     # run_on_click=True,
                 #     # label="Select an example image to test"
                 # )
-        gallery.select(load_img, outputs=[input_image, original_image, selected_points])
+        gallery.select(load_img, outputs=[input_image, original_image, original_image_high_res, selected_points])
 
-        input_image.upload(store_img, [input_image], [input_image, original_image, selected_points])
+        input_image.upload(store_img, [input_image], [input_image, original_image, original_image_high_res, selected_points])
 
         # Get points and draw arrows or zero-length vectors based on the toggle
         def get_point(img, sel_pix, zero_length, evt: gr.SelectData):
@@ -251,7 +252,7 @@ with gr.Blocks() as demo:
         def run_model_on_points(points, input_image, original_image):
             H = input_image.shape[0]
             W = input_image.shape[1]
-            factor = 224/H
+            factor = 256/H
             # Example: pretend the model processes points and returns a simple transformation on the image
             points = torch.from_numpy(np.array(points).reshape(-1, 4)) * factor
 
@@ -260,8 +261,8 @@ with gr.Blocks() as demo:
             img = Image.fromarray(original_image)
 
             transform = transforms.Compose([
-                transforms.Resize((224, 224)),
-                transforms.CenterCrop(224)
+                transforms.Resize((256, 256)),
+                # transforms.CenterCrop(256)
             ])
             img = np.array(transform(img))
 
@@ -272,7 +273,7 @@ with gr.Blocks() as demo:
             img = img[None]
 
             # reshape image to [B, C, T, H, W], C = 3, T = 3 (3-frame model), H = W = 224
-            x = img[:, :, None].expand(-1, -1, 3, -1, -1)#.to(torch.float16)
+            x = img[:, :, None].expand(-1, -1, 2, -1, -1)#.to(torch.float16)
 
             # Imagenet-normalize the inputs (standardization)
 
@@ -290,9 +291,9 @@ with gr.Blocks() as demo:
             return counterfactual
 
         # Run model when the button is clicked
-        run_model_button.click(run_model_on_points, [selected_points, input_image, original_image], [output_image])
+        run_model_button.click(run_model_on_points, [selected_points, input_image, original_image_high_res], [output_image])
 
 
 
     # Launch the app
-demo.queue().launch()
+demo.queue().launch(inbrowser=True, share=True)

cwm/model/model_factory.py

@@ -26,6 +26,11 @@ _model_catalogue ={
         "init_fn": model_pretrain.vitb_8x8patch_2frames,
     },
 
+    "vitb_8x8patch_2frames_encoder_mask_token": {
+        "path": "cwm/2frame_cwm_mask_token.pth",
+        "init_fn": model_pretrain.vitb_8x8patch_2frames_encoder_mask_token,
+    },
+
 }
 
 

cwm/model/model_pretrain.py

@@ -818,6 +818,23 @@ def vitb_4x4patch_2frames(**kwargs):
         **kwargs)
     return model
 
+from cwm.model.modeling_pretrain_cleaned_soft import pretrain_vit_base_256_scaffold
+
+def vitb_8x8patch_2frames_encoder_mask_token(
+        use_flash_attention=False, **kwargs):
+    model = pretrain_vit_base_256_scaffold(
+        patch_size=(8, 8),
+        num_frames=2,
+        tubelet_size=1,
+        use_flash_attention=use_flash_attention,
+        interp_noise=False,
+        legacy=False,
+        xla_flash=False,
+        learn_pos_embed=True,
+        **kwargs)
+    return model
+
+
 # def base_8x8patch_2frames_1tube(**kwargs):
 #     model = pretrain_videomae_base_224_scaffold(
 #         patch_size=(8, 8),

cwm/model/modeling_pretrain_cleaned_soft.py

@@ -0,0 +1,723 @@
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import trunc_normal_ as __call_trunc_normal_
+from einops import rearrange
+from cwm.model.model_utils import Block, _cfg, PatchEmbed, get_sinusoid_encoding_table
+
+from torch import Tensor
+import cwm.utils as utils
+
+def trunc_normal_(tensor, mean=0., std=1.):
+    __call_trunc_normal_(tensor, mean=mean, std=std, a=-std, b=std)
+
+
+def interpolate_pos_encoding(pos_embed, n_frames, h, w):
+    N = pos_embed.shape[1]
+    if N == (h * w * n_frames):
+        return pos_embed
+    old_h = old_w = int((N / n_frames) ** 0.5)
+    patch_pos_embed = pos_embed.view(1, n_frames, old_h, old_w, -1).flatten(0, 1).permute(0, 3, 1, 2)
+
+    patch_pos_embed = F.interpolate(
+        patch_pos_embed,
+        size=(h, w),
+        mode='bicubic',
+    )
+    return patch_pos_embed.permute(0, 2, 3, 1).flatten(0, 2).unsqueeze(0)
+
+PRINT_PADDING = False
+
+class PretrainVisionTransformerEncoder(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self, img_size=224, patch_size=(16, 16), in_chans=3, num_classes=0, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None, tubelet_size=2,
+                 use_learnable_pos_emb=False, num_frames=16, embed_per_frame=False, clumping_factor=None, block_func=Block, k_bias=False, interp_noise=False, block_kwargs={}, legacy=False, xla_flash=False, learn_pos_embed=False):
+        super().__init__()
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.patch_size = (tubelet_size,) + patch_size
+        self.pt, self.ph, self.pw = self.patch_size
+        self.h = int(img_size / self.ph)
+        self.w = int(img_size / self.pw)
+        self.hw = self.h * self.w
+
+        self.clumping_factor = clumping_factor
+        self.interp_noise = interp_noise
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+
+        if self.clumping_factor is not None:  # Clump the context frame for memory efficiency
+            self.clumping_embed = nn.Conv3d(in_channels=embed_dim, out_channels=embed_dim,
+                                            kernel_size=(tubelet_size, clumping_factor, clumping_factor),
+                                            stride=(tubelet_size, clumping_factor, clumping_factor))
+
+        self._embed_per_frame = embed_per_frame
+        if not self._embed_per_frame:
+            self.patch_embed = PatchEmbed(
+                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,tubelet_size=tubelet_size,num_frames=num_frames)
+            num_patches = self.patch_embed.num_patches
+        elif self._embed_per_frame:
+            assert (num_frames % tubelet_size) == 0
+            num_embeddings = (num_frames // tubelet_size)
+            self.patch_embed = nn.ModuleList([
+                PatchEmbed(
+                    img_size=img_size, patch_size=patch_size,
+                    in_chans=in_chans, embed_dim=embed_dim,
+                    tubelet_size=tubelet_size, num_frames=tubelet_size)
+                for _ in range(num_embeddings)])
+            num_patches = self.patch_embed[0].num_patches * num_embeddings
+
+        self.num_patches = num_patches
+        self.num_frames = num_frames
+        print("NUM PATCHES IN ENCODER", self.num_patches)
+
+        self.pos_embed = get_sinusoid_encoding_table(num_patches, embed_dim)
+
+        if learn_pos_embed:
+            self.pos_embed = nn.Parameter(self.pos_embed)
+
+        self.learn_pos_embed = learn_pos_embed
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            block_func(
+                dim=embed_dim, in_dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                init_values=init_values, **block_kwargs, k_bias=k_bias, legacy=legacy, xla_flash=xla_flash)
+            for i in range(depth)])
+        self.norm =  norm_layer(embed_dim)
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        if use_learnable_pos_emb:
+            trunc_normal_(self.pos_embed, std=.02)
+
+        self.apply(self._init_weights)
+
+    def _set_pos_embed(self, dim=None):
+        if dim is None:
+            dim = self.embed_dim
+        if self.pos_embed is None:
+            self.pos_embed = get_sinusoid_encoding_table(
+                self.num_patches, dim)
+
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def _get_pos_embed(self):
+        return self.pos_embed
+
+    def forward_block(self, x, idx):
+        return self.blocks[idx](x)
+
+    def interpolate_tensor_with_mask_token(self,
+            x: Tensor, mask: Tensor, mask_token: Tensor, invert: bool = True
+    ) -> Tensor:
+        """
+        Where mask == (0 if invert else 1), return x
+        where mask == (1 if invert else 0), return mask_token
+        Linearly interpolate between these using value of mask.
+        """
+        # mask_token = mask_token
+        # breakpoint()
+        B, N, C = x.shape
+        assert mask.shape[1] == N, (
+            f"Number of tokens in mask ({mask.shape[1]}) does not match "
+            f"number of tokens in input ({N})"
+        )
+
+        assert mask_token.shape[-1] == C, (
+            f"Dimensionality of mask token ({mask_token.shape[-1]}) does not match "
+            f"dimensionality of tokens in input ({C})"
+        )
+
+        # convert mask to interpolation weights in range [0., 1.]
+        mask = mask.to(x).clip(min=0.0, max=1.0)
+        mask = (1.0 - mask) if invert else mask
+        mask = mask.unsqueeze(-1)  # [B, N, 1]
+
+        # expand mask token
+        mask_token = mask_token.view(1, 1, C).expand(B, N, -1)
+
+        # interpolate
+        start = mask_token
+        end = x
+
+        return start + mask * (end - start)
+
+    def interpolate_tensor_with_noise(self,
+            x: Tensor, mask: Tensor, invert: bool = True
+    ) -> Tensor:
+        """
+        Where mask == (0 if invert else 1), return x
+        where mask == (1 if invert else 0), return mask_token
+        Linearly interpolate between these using value of mask.
+        """
+        # mask_token = mask_token
+        # breakpoint()
+        B, N, C = x.shape
+        assert mask.shape[1] == N, (
+            f"Number of tokens in mask ({mask.shape[1]}) does not match "
+            f"number of tokens in input ({N})"
+        )
+
+        # convert mask to interpolation weights in range [0., 1.]
+        mask = mask.to(x).clip(min=0.0, max=1.0)
+        mask = (1.0 - mask) if invert else mask
+        mask = mask.unsqueeze(-1)  # [B, N, 1]
+
+        # ImageNet mean and std
+        mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+        std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+
+        # Generate a 3x8x8 patch of random numbers from a normal distribution
+        # with the same mean and std as ImageNet images
+        rand_vec = torch.randn(B, N, 3, self.patch_size[-2], self.patch_size[-1]) * std + mean
+
+        rand_vec = rand_vec.to(x.device).to(x.dtype).view(B, N, -1)
+        # interpolate
+        start = rand_vec
+        end = x
+
+        return start + mask * (end - start)
+
+    def tokenize(self, x, mask=None):
+
+        if not self._embed_per_frame:
+            x = self.patch_embed(x)
+        elif self._embed_per_frame:
+            x = torch.cat([
+                self.patch_embed[i](
+                    x[:,:,(i*self.pt):((i+1)*self.pt)])
+                for i in range(len(self.patch_embed))], 1)
+            
+        pos_embed = self._get_pos_embed().type_as(x).to(x.device).clone()
+        if not self._learnable_pos_embed:
+            pos_embed = pos_embed.detach()
+        x = x + pos_embed
+        return (x, mask)
+
+    def tokenize_and_mask(self, x, mask):
+
+        x, mask = self.tokenize(x, mask)
+        B, _, C = x.shape
+        # breakpoint()
+        x_vis = x[~mask].reshape(B, -1, C)
+        return x_vis
+
+    def tokenize_and_mask_variable_size(self, x, mask):
+
+        x, mask = self.tokenize(x, mask)
+        B, _, C = x.shape
+        all_batches = []
+        max_len = 0
+        all_len = []
+        for i in range(B):
+            x_vis = x[i, ~mask[i]]
+            if x_vis.shape[0] > max_len:
+                max_len = x_vis.shape[0]
+            all_batches.append(x_vis)
+            all_len.append(x_vis.shape[0])
+
+        #pad all batches to max_len in a single line
+        x_vis = torch.stack([F.pad(batch, (0,0,0,max_len-batch.shape[0]), mode='constant', value=0) for batch in all_batches])
+
+        return x_vis, all_len
+
+    def forward_features(self, x, mask, move_patches, static_patches, delta, mask_token, res=1, return_feat_layer=None):
+        _, _, T, H, W = x.shape
+
+        if self.interp_noise:
+            #patchify x with patch size[0], patch size[1]
+            p0 = self.patch_size[-2]
+            p1 = self.patch_size[-1]
+            x = rearrange(x, 'b c t (h p0) (w p1) -> b (t h w) (p0 p1 c)', p0=p0, p1=p1, h=H//p0, w=W//p1) # x: [B, N, C]
+
+            x = self.interpolate_tensor_with_noise(x, mask, invert=True)
+            x = rearrange(x, 'b n (p c) -> b n p c', c=3)
+            # Notice: To visualize the reconstruction video, we add the predict and the original mean and var of each patch.
+            x = rearrange(x,
+                            'b (t h w) (p0 p1 p2) c -> b c (t p0) (h p1) (w p2)',
+                            p0=1,
+                            p1=self.patch_size[-2],
+                            p2=self.patch_size[-1],
+                            h=H//self.patch_size[-2],
+                            w=W//self.patch_size[-1])
+
+        x = embed = self.patch_embed(x)
+
+        if res != 1:
+           
+            p0 = self.patch_size[-2]
+            p1 = self.patch_size[-1]
+            pos_embed = interpolate_pos_encoding(self.pos_embed, T, int(256 // p0 * res), int(256 // p1 * res))
+        else:
+           
+            pos_embed = self._get_pos_embed()
+
+        pos_embed = pos_embed.type_as(x)  # .to(x.device).clone()
+
+        if not self.learn_pos_embed:
+            pos_embed = pos_embed.to(x.device).clone().detach()
+
+        x = x + pos_embed
+        B, _, C = x.shape
+        # x_vis = x[~mask].reshape(B, -1, C) # ~mask means visible
+        if not self.interp_noise:
+            x_vis = self.interpolate_tensor_with_mask_token(x, mask, mask_token, invert=True)
+        else:
+            x_vis = x
+
+        if move_patches is not None:
+
+            assert B == 1, "Only support batch size 1 for now"
+            for (px, py) in move_patches:
+                idx = px * self.w + py
+                dx, dy = delta
+                nx, ny = px + dx, py + dy
+                new_idx = nx * self.w + ny + (self.patch_embed.num_frames - 1) * (self.h * self.w)
+
+                emb = embed[:, idx]
+                pos_emb = pos_embed[:, new_idx]
+                emb = emb + pos_emb
+                x_vis = torch.cat([x_vis, emb[None]], 1)
+
+            if static_patches is not None:
+                for (px, py) in static_patches:
+                    idx = px * self.w + py
+                    new_idx = px * self.w + py + (self.patch_embed.num_frames - 1) * (self.h * self.w)
+                    emb = embed[:, idx]
+                    pos_emb = pos_embed[:, new_idx]
+                    emb = emb + pos_emb
+                    x_vis = torch.cat([x_vis, emb[None]], 1)
+
+        for blk_idx, blk in enumerate(self.blocks):
+            x_vis = blk(x_vis)
+            if blk_idx == return_feat_layer:
+                return x_vis
+
+        x_vis = self.norm(x_vis)
+        return x_vis
+
+    def _set_inputs(self, *args, **kwargs):
+        pass
+
+    def forward(self, x, mask, mask_token, return_feat_layer=None, timestamps=None, move_patches=None, static_patches=None, delta=None, res=1):
+        self._set_inputs(x, mask)
+        # pass input through the encoder
+        x = self.forward_features(x, mask, move_patches, static_patches, delta, mask_token, return_feat_layer=return_feat_layer, res=res)
+        # if return_feat_layer is not None and is lesser than the number of blocks it means that we are returning the
+        # features of an intermediate block layer. in this case we do not want to apply the head layer
+        if return_feat_layer is not None and return_feat_layer < len(self.blocks):
+            return x
+        # if we are passing through the entire encoder transformer we apply the head layer
+        x = self.head(x)
+        return x
+
+class PretrainVisionTransformerDecoder(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self, patch_size=(16, 16), num_classes=768, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None, block_func=Block, block_kwargs={}, k_bias=False, legacy=True, xla_flash=False
+                 ):
+        super().__init__()
+
+
+        self.num_classes = num_classes
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.patch_size = patch_size
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            block_func(
+                dim=embed_dim, in_dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                init_values=init_values, **block_kwargs, k_bias=k_bias, legacy=legacy, xla_flash=xla_flash)
+            for i in range(depth)])
+        self.norm =  norm_layer(embed_dim)
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.apply(self._init_weights)
+
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_block(self, x, idx):
+        return self.blocks[idx](x)
+
+    def get_last_tokens(self, x, return_token_num):
+        if return_token_num > 0:
+            return self.head(self.norm(x[:,-return_token_num:]))
+        elif return_token_num == 0:
+            return self.head(self.norm(x))[:,x.size(1):]
+        else:
+            return self.head(self.norm(x))
+
+    def forward(self, x, return_token_num, return_feat_layer=None):
+
+        # pass input through the decoder        
+        for blk_idx, blk in enumerate(self.blocks):
+            x = blk(x)
+            # if we are returning the features of an intermediate block
+            # do so and skip the remaining computation
+            if blk_idx == return_feat_layer:
+                return x
+
+        if return_token_num > 0:
+            x = self.head(self.norm(x[:, -return_token_num:])) # only return the mask tokens predict pixels
+        else:
+            x = self.head(self.norm(x))
+
+        return x
+
+class PretrainVisionTransformer(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    default_input_kwargs = {'unnormalize': True}
+    def __init__(self,
+                 img_size=224, 
+                 patch_size=(16, 16),
+                 main_input=None,
+                 main_input_kwargs=default_input_kwargs,
+                 encoder_func=PretrainVisionTransformerEncoder,
+                 encoder_in_chans=3, 
+                 encoder_num_classes=0, 
+                 encoder_embed_dim=768, 
+                 encoder_depth=12,
+                 encoder_num_heads=12,
+                 encoder_block_func=Block,
+                 encoder_block_kwargs={},
+                 decoder_num_classes=None, # For pretraining this parameter isn't relevant but must be set according to tube&patch size
+                 decoder_embed_dim=512, 
+                 decoder_depth=8,
+                 decoder_num_heads=8,
+                 decoder_block_func=Block,
+                 decoder_block_kwargs={},
+                 mlp_ratio=4., 
+                 qkv_bias=False,
+                 k_bias=False,
+                 qk_scale=None,
+                 num_frames=16,
+                 drop_rate=0., 
+                 attn_drop_rate=0.,
+                 drop_path_rate=0., 
+                 norm_layer=nn.LayerNorm, 
+                 init_values=0.,
+                 spacetime_separable_pos_embed=False,
+                 tubelet_size=2,
+                 num_classes=0, # avoid the error from create_fn in timm
+                 in_chans=0, # avoid the error from create_fn in timm
+                 embed_per_frame=False,
+                 flow_model_ckpt=None,
+                 flow_frames=None,
+                 random_input=False,
+                 use_flash_attention=False,
+                 additional_decoder_for_transition=False,
+                 additional_decoder_for_x3_hat=False,
+                 clumping_factor=None,
+                 return_detectron_format=False,
+                 out_feature='out_feature',
+                 interp_noise=False,
+                 legacy=True,
+                 xla_flash=False,
+                 learn_pos_embed=False,
+                 **kwargs
+                 ):
+        super().__init__()
+
+        encoder_block_kwargs.update({'flash_attention': use_flash_attention})
+        decoder_block_kwargs.update({'flash_attention': use_flash_attention})
+
+        self.clumping_factor = clumping_factor
+
+        self.interp_noise = interp_noise
+
+        self.learn_pos_embed = learn_pos_embed
+
+        if self.clumping_factor is not None:
+            print('Clumping factor = %d' % self.clumping_factor)
+            self.clumping_embed = nn.Conv3d(in_channels=decoder_embed_dim, out_channels=decoder_embed_dim,
+                                            kernel_size=(1, clumping_factor, clumping_factor),
+                                            stride=(1, clumping_factor, clumping_factor))
+            self.clumping_embed.apply(self._init_weights)
+
+            self.up = nn.ConvTranspose2d(decoder_embed_dim, decoder_embed_dim, kernel_size=2, stride=2)
+            self.up.apply(self._init_weights)
+
+        self.encoder = encoder_func(
+            img_size=img_size, 
+            patch_size=patch_size, 
+            in_chans=encoder_in_chans, 
+            num_classes=encoder_num_classes, 
+            embed_dim=encoder_embed_dim, 
+            depth=encoder_depth,
+            num_heads=encoder_num_heads, 
+            mlp_ratio=mlp_ratio, 
+            qkv_bias=qkv_bias, 
+            qk_scale=qk_scale, 
+            drop_rate=drop_rate, 
+            attn_drop_rate=attn_drop_rate,
+            drop_path_rate=drop_path_rate, 
+            norm_layer=norm_layer, 
+            init_values=init_values,
+            tubelet_size=tubelet_size,
+            num_frames=num_frames,
+            embed_per_frame=embed_per_frame,
+            block_func=encoder_block_func,
+            block_kwargs=encoder_block_kwargs,
+            clumping_factor=clumping_factor,
+            k_bias=k_bias,
+            interp_noise = interp_noise,
+            legacy=legacy,
+            xla_flash=xla_flash,
+            learn_pos_embed=learn_pos_embed,
+            **kwargs)
+
+        if not return_detectron_format:
+            self.decoder = PretrainVisionTransformerDecoder(
+                patch_size=patch_size,
+                num_classes= 3*tubelet_size*(patch_size[0]*patch_size[1]) if decoder_num_classes is None else decoder_num_classes,
+                embed_dim=decoder_embed_dim,
+                depth=decoder_depth,
+                num_heads=decoder_num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop_rate=drop_rate,
+                attn_drop_rate=attn_drop_rate,
+                drop_path_rate=drop_path_rate,
+                norm_layer=norm_layer,
+                init_values=init_values,
+                block_func=decoder_block_func,
+                k_bias=k_bias, xla_flash=xla_flash,
+                block_kwargs=decoder_block_kwargs, legacy=legacy)
+
+            self.encoder_to_decoder = nn.Linear(encoder_embed_dim, decoder_embed_dim, bias=k_bias)
+
+            if not self.interp_noise:
+                self.mask_token = nn.Parameter(torch.zeros(1, 1, encoder_embed_dim))
+                trunc_normal_(self.mask_token, std=.02)
+            else:
+                self.mask_token = None
+
+        self.timestamps = None
+        self.encoder.timestamps = None
+
+        if self.learn_pos_embed:
+            self.pos_embed = nn.Parameter(get_sinusoid_encoding_table(self.encoder.num_patches, decoder_embed_dim))
+        else:
+            self.pos_embed = get_sinusoid_encoding_table(self.encoder.num_patches, decoder_embed_dim)
+
+        self.num_frames = num_frames
+        self.num_patches = self.encoder.num_patches
+        if self.num_frames is not None:
+            self.num_patches_per_frame = self.num_patches // self.num_frames
+        else:
+            self.num_patches_per_frame = self.num_patches
+        self.patch_size = self.encoder.patch_size
+        if isinstance(img_size, int):
+            self.image_size = (img_size, img_size)
+        else:
+            assert hasattr(img_size, '__len__'), img_size
+            self.image_size = img_size
+
+        self.return_detectron_format = return_detectron_format
+
+    @property
+    def mask_size(self):
+        return (self.num_frames // self.patch_size[0],
+                self.image_size[-2] // self.patch_size[-2],
+                self.image_size[-1] // self.patch_size[-1])
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token', 'mask_token'}
+
+
+
+    def unpatchify(self, x, mask):
+        # Define the input tensor
+        B, N, C = x.shape  # batch size
+        h, w = self.mask_size[-2:]
+        patch_size = self.patch_size[-2:]
+
+        recon = torch.zeros(B, h*w, C).to(x)
+        recon[mask[:, -h*w:]] = x.flatten(0, 1)
+
+        rec_imgs = rearrange(recon, 'b n (p c) -> b n p c', c=3)
+        # Notice: To visualize the reconstruction video, we add the predict and the original mean and var of each patch.
+        rec_imgs = rearrange(rec_imgs,
+                             'b (t h w) (p0 p1 p2) c -> b c (t p0) (h p1) (w p2)',
+                             p0=1,
+                             p1=patch_size[0],
+                             p2=patch_size[1],
+                             h=h,
+                             w=w)
+
+        # MEAN = torch.from_numpy(np.array((0.485, 0.456, 0.406))[None, :, None, None, None]).cuda().half()
+        # STD = torch.from_numpy(np.array((0.229, 0.224, 0.225))[None, :, None, None, None]).cuda().half()
+        #
+        # rec_imgs = (rec_imgs - MEAN) / STD
+
+        return rec_imgs
+
+
+    def forward(self, x, mask, timestamps=None, return_feat_layer=None, res=1, *args, get_encoder_out=False, **kwargs):
+
+        _, _, T, _, _ = x.shape
+
+        self.device = x.device
+
+        enc_out = self.encoder(x, mask, self.mask_token, timestamps=timestamps, return_feat_layer=return_feat_layer, res=res, *args, **kwargs) # [B, N_vis, C_e]
+
+        x_vis = self.encoder_to_decoder(enc_out)
+
+        # check if we are returning the features of an intermediate block layer
+        if return_feat_layer is not None:
+            # if the returned layer is one of the encoder layers (the first N_enc layers) we return the features
+            # if the return feat layer is exactly N_enc then we are returning the layer after the entire encoder block
+            # in both cases this manifests as returning x_vis, since self.encoder will return either the final block embedding
+            # or the final head embedding depending on the return_feat_layer
+            # in either case we subtract the number of encoder blocks + 1 (for the intermediate embedding layer)
+            # from the return_feat_layer to get the correct index for the decoder block
+            return_feat_layer = return_feat_layer - len(self.encoder.blocks) - 1
+            if return_feat_layer < 0:
+                return x_vis
+
+        # add pos embedding
+        if res != 1:
+            p0 = self.patch_size[-2]
+            p1 = self.patch_size[-1]
+            pos_embed = interpolate_pos_encoding(self.pos_embed, T, int(256 // p0 * res), int(256 // p1 * res))
+        else:
+            pos_embed = self.pos_embed
+        dec_pos_embed = pos_embed.expand(x_vis.size(0), -1, -1).type_as(x)
+
+        if not self.learn_pos_embed:
+            dec_pos_embed = dec_pos_embed.to(x.device).clone().detach()
+
+        x_vis = x_vis + dec_pos_embed
+
+        # pass input through the decoder, this will automatically return an intermediate layer if return_feat_layer is set
+        x_all = self.decoder(x_vis, 0, return_feat_layer=return_feat_layer)
+
+        if get_encoder_out:
+            return x_all, enc_out
+        
+        return x_all
+
+    def get_counterfactual(self, x, move_patches):
+        '''
+        :param x: input tensor [1, C, T, H, W]: support only batch size 1 for now
+        :param move_patches: torch tensor [N, 4] sized array where each row contains patch motion [x1, y1, x2, y2] in pixel coordinates
+        :return:
+        '''
+        B, _, T, H, H = x.shape
+
+        mask = torch.ones(B, self.encoder.hw * self.encoder.num_frames).to(x.device).bool()
+        mask[:, :self.encoder.hw * (self.encoder.num_frames - 1)] = False
+
+        move_patches = (move_patches / H) * self.encoder.h
+        move_patches = move_patches.to(torch.int64)
+
+        for x1, y1, x2, y2 in move_patches:
+            idx2 = x2 * self.encoder.w + y2 + (self.encoder.num_frames - 1) * (self.encoder.h * self.encoder.w)
+            mask[:, idx2] = False
+            im_x1 = x1 * self.encoder.ph
+            im_y1 = y1 * self.encoder.pw
+            im_x2 = x2 * self.encoder.ph
+            im_y2 = y2 * self.encoder.pw
+            x[:, :, -1, im_x2:im_x2 + self.encoder.ph, im_y2:im_y2 + self.encoder.pw] = x[:, :, -2,
+                                                                                        im_x1:im_x1 + self.encoder.ph,
+                                                                                        im_y1:im_y1 + self.encoder.pw]
+
+        prediction = self.forward(x, mask)[:, -self.encoder.hw:]
+
+        prediction = utils.unpatchify_cwm(
+            prediction,
+            patch_size=self.encoder.patch_size[-1],
+        )  # reshape the output to an image
+
+        return prediction
+
+
+def pretrain_vit_base_256_scaffold(**kwargs):
+    model = PretrainVisionTransformer(
+        img_size=256,
+        encoder_embed_dim=768,
+        encoder_depth=12,
+        encoder_num_heads=12,
+        encoder_num_classes=0,
+        decoder_embed_dim=768,
+        decoder_num_heads=12,
+        decoder_depth=12,
+        mlp_ratio=4,
+        
+        qkv_bias=True,
+        k_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs)
+    model.default_cfg = _cfg()
+    return model
+