cwm/model/modeling_pretrain_cleaned_soft.py

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