vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/UFM/UniCeption/uniception/models/encoders/dense_rep_encoder.py

@@ -0,0 +1,344 @@
+"""
+Encoder class for Dense Representation Encoder
+"""
+
+import math
+from functools import partial
+from typing import Callable, List, Optional, Tuple, Type, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.init import trunc_normal_
+
+from uniception.models.encoders.base import (
+    UniCeptionViTEncoderBase,
+    ViTEncoderInput,
+    ViTEncoderNonImageInput,
+    ViTEncoderOutput,
+)
+
+
+def make_2tuple(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class ResidualBlock(nn.Module):
+    "Redidual block for Dense Representation Encoder"
+
+    def __init__(self, in_channels: int, out_channels: int, act_layer: Type[nn.Module] = nn.GELU):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.act = act_layer()
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.shortcut = (
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+
+    def forward(self, x):
+        identity = self.shortcut(x)
+        out = self.conv1(x)
+        out = self.act(out)
+        out = self.conv2(out)
+        out += identity
+
+        return self.act(out)
+
+
+class DenseRepresentationEncoder(UniCeptionViTEncoderBase):
+    "UniCeption Dense Representation Encoder"
+
+    def __init__(
+        self,
+        name: str,
+        in_chans: int = 3,
+        enc_embed_dim: int = 1024,
+        apply_pe: bool = True,
+        input_size_for_pe: Union[int, Tuple[int, int]] = 518,
+        patch_size: int = 14,
+        intermediate_dims: List[int] = [588, 768, 1024],
+        data_norm_type: str = "dense_rep_encoder",
+        act_layer: Type[nn.Module] = nn.GELU,
+        norm_layer: Optional[Callable] = partial(nn.LayerNorm, eps=1e-6),
+        post_pe_norm_layer: Optional[Callable] = partial(nn.LayerNorm, eps=1e-6),
+        interpolate_antialias: bool = False,
+        interpolate_offset: float = 0.1,
+        pretrained_checkpoint_path: str = None,
+        *args,
+        **kwargs,
+    ):
+        """
+        Dense Representation Encoder for extracting patch-wise features from a spatial input of size (B, C, H, W).
+        Uses a convolution based patchify followed by some residual blocks.
+        Also applies positional encoding with interpolation to the patch-wise features if required.
+
+        Args:
+            in_chans (int): Number of input channels.
+            enc_embed_dim (int): Embedding dimension of the encoder.
+            apply_pe (bool): Whether to apply positional encoding.
+            input_size_for_pe (Union[int, Tuple[int, int]]): Input size for positional encoding.
+            patch_size (int): Patch size of the encoder.
+            intermediate_dims (List[int]): Intermediate dimensions of the encoder.
+            data_norm_type (str): Data normalization type. (Used for checking if the input images are normalized correctly.)
+            act_layer (Type[nn.Module]): Activation layer.
+            norm_layer (Optional[Callable]): Normalization layer.
+            post_pe_norm_layer (Optional[Callable]): Normalization layer after positional encoding.
+            interpolate_antialias (bool): Whether to apply antialiasing in interpolation.
+            interpolate_offset (float): Offset for interpolation.
+            pretrained_checkpoint_path (str): Path to pretrained checkpoint.
+        """
+        # Init the base class
+        super().__init__(
+            name=name,
+            data_norm_type=data_norm_type,
+            patch_size=patch_size,
+            *args,
+            **kwargs,
+        )
+
+        # Init the specific attributes
+        self.in_chans = in_chans
+        self.enc_embed_dim = enc_embed_dim
+        self.intermediate_dims = intermediate_dims
+        self.apply_pe = apply_pe
+
+        # Initialize the encoder with a pixel unshuffle and conv projection to patchify the input
+        self.unshuffle = nn.PixelUnshuffle(self.patch_size)
+        self.conv_in = nn.Conv2d(self.in_chans * (self.patch_size**2), self.intermediate_dims[0], 3, 1, 1)
+
+        # Add residual blocks
+        layers = []
+        for intermediate_idx in range(len(self.intermediate_dims) - 1):
+            layers.append(
+                ResidualBlock(
+                    in_channels=self.intermediate_dims[intermediate_idx],
+                    out_channels=self.intermediate_dims[intermediate_idx + 1],
+                    act_layer=act_layer,
+                )
+            )
+
+        # Final projection to match encoder embeddings dim
+        layers.append(
+            nn.Conv2d(
+                in_channels=self.intermediate_dims[-1],
+                out_channels=self.enc_embed_dim,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+            )
+        )
+        self.encoder = nn.Sequential(*layers)
+
+        # Init norm layer after encoder if required
+        self.norm_layer = norm_layer(enc_embed_dim) if norm_layer else nn.Identity()
+        if isinstance(self.norm_layer, nn.LayerNorm):
+            nn.init.constant_(self.norm_layer.bias, 0)
+            nn.init.constant_(self.norm_layer.weight, 1.0)
+
+        if self.apply_pe:
+            # Init the patch resolution details required for positional encoding
+            patch_HW = make_2tuple(patch_size)
+            self.input_size_for_pe = make_2tuple(input_size_for_pe)
+            self.patches_resolution = (
+                self.input_size_for_pe[0] // patch_HW[0],
+                self.input_size_for_pe[1] // patch_HW[1],
+            )
+            self.num_patches = self.patches_resolution[0] * self.patches_resolution[1]
+
+            # Init the sinusodial positional encodings
+            self.register_buffer(
+                "pos_embed",
+                self._get_sinusoid_encoding_table(self.num_patches, self.enc_embed_dim, 70007),
+            )
+            self.interpolate_antialias = interpolate_antialias
+            self.interpolate_offset = interpolate_offset
+
+            # Init the norm layer after positional encoding if required
+            self.post_pe_norm = post_pe_norm_layer(enc_embed_dim) if post_pe_norm_layer else nn.Identity()
+            if isinstance(self.post_pe_norm, nn.LayerNorm):
+                nn.init.constant_(self.post_pe_norm.bias, 0)
+                nn.init.constant_(self.post_pe_norm.weight, 1.0)
+
+        # Load the pretrained checkpoint if provided
+        self.pretrained_checkpoint_path = pretrained_checkpoint_path
+        if self.pretrained_checkpoint_path:
+            print(
+                f"Loading custom pretrained Dense Representation Encoder checkpoint from {self.pretrained_checkpoint_path} ..."
+            )
+            ckpt = torch.load(self.pretrained_checkpoint_path, weights_only=False)
+            print(self.load_state_dict(ckpt["model"]))
+
+    def _get_sinusoid_encoding_table(self, n_position, d_hid, base):
+        "Sinusoid position encoding table"
+
+        def get_position_angle_vec(position):
+            return [position / np.power(base, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]
+
+        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
+        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])
+        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])
+
+        return torch.FloatTensor(sinusoid_table)
+
+    def interpolate_pos_encoding(self, features, height, width):
+        """
+        Interpolate the positional encoding to the expected size.
+
+        Args:
+            features (torch.Tensor): Input tensor of shape (B, N, C).
+            height (int, float): Height of the input tensor.
+            width (int, float): Width of the input tensor.
+
+        Returns:
+            torch.Tensor: Interpolated positional encoding tensor of shape (1, N, C).
+        """
+        previous_dtype = features.dtype
+        npatch = features.shape[1]
+        N = self.pos_embed.unsqueeze(0).shape[1]
+        if npatch == N and height == width:
+            return self.pos_embed.unsqueeze(0)
+        patch_pos_embed = self.pos_embed.unsqueeze(0).float()
+        dim = features.shape[-1]
+        height0 = height // self.patch_size
+        width0 = width // self.patch_size
+        M = int(math.sqrt(N))  # Recover the number of patches in each dimension
+        assert N == M * M
+        kwargs = {}
+        if self.interpolate_offset:
+            # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8
+            # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors
+            sh = float(height0 + self.interpolate_offset) / M
+            sw = float(width0 + self.interpolate_offset) / M
+            kwargs["scale_factor"] = (sh, sw)
+        else:
+            # Simply specify an output size instead of a scale factor
+            kwargs["size"] = (height0, width0)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2),
+            mode="bicubic",
+            antialias=self.interpolate_antialias,
+            **kwargs,
+        )
+        assert (height0, width0) == patch_pos_embed.shape[-2:]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+        return patch_pos_embed.to(previous_dtype)
+
+    def forward(self, encoder_input: Union[ViTEncoderInput, ViTEncoderNonImageInput]) -> ViTEncoderOutput:
+        """
+        Dense Representation Encoder Forward Pass
+
+        Args:
+            encoder_input (Union[ViTEncoderInput, ViTEncoderNonImageInput]): Input data for the encoder.
+                If input type is ViTEncoderInput, input data must contain image normalization type and normalized image tensor.
+                If input type is ViTEncoderNonImageInput, input data must contain a tensor of size (B, C, H, W).
+
+        Returns:
+            ViTEncoderOutput: Output data from the encoder.
+        """
+        # Get the input data and verify normalization if the input type is ViTEncoderInput
+        if isinstance(encoder_input, ViTEncoderInput):
+            self._check_data_normalization_type(encoder_input.data_norm_type)
+            input_data = encoder_input.image
+        elif isinstance(encoder_input, ViTEncoderNonImageInput):
+            input_data = encoder_input.data
+        else:
+            raise ValueError("Unsupported input type for Dense Representation Encoder.")
+
+        # Check the dtype and shape of the input
+        assert isinstance(input_data, torch.Tensor), "Input must be a torch.Tensor"
+        assert input_data.ndim == 4, "Input must be of shape (B, C, H, W)"
+        assert input_data.shape[1] == self.in_chans, f"Input channels must be {self.in_chans}"
+        batch_size, channels, height, width = input_data.shape
+        assert (
+            height % self.patch_size == 0 and width % self.patch_size == 0
+        ), f"Input shape must be divisible by patch size: {self.patch_size}"
+
+        # Encode the dense representation
+        features = self.unshuffle(input_data)
+        features = self.conv_in(features)
+        features = self.encoder(features)
+        features = features.flatten(2).transpose(
+            1, 2
+        )  # (B, E, H / Patch_Size, W / Patch_Size) -> (B, H / Patch_Size * W / Patch_Size, E)
+        features = self.norm_layer(features)  # Normalize the features after patch encoding
+
+        # Apply positional encoding if required
+        if self.apply_pe:
+            features = features + self.interpolate_pos_encoding(
+                features, height, width
+            )  # (B, H / Patch_Size * W / Patch_Size, E)
+            features = self.post_pe_norm(features)  # Normalize the features after positional encoding
+
+        # Resize the features to the expected shape
+        # (B x Num_patches x Embed_dim) -> (B x Embed_dim x H / Patch_Size x W / Patch_Size)
+        features = features.permute(0, 2, 1)
+        features = features.reshape(
+            -1, self.enc_embed_dim, height // self.patch_size, width // self.patch_size
+        ).contiguous()
+
+        return ViTEncoderOutput(features=features)
+
+
+if __name__ == "__main__":
+    # Init Dense Representation Encoder for images as input
+    patch_embedder = DenseRepresentationEncoder(
+        name="dense_rep_encoder",
+        data_norm_type="dense_rep_encoder",
+        input_size_for_pe=518,
+        patch_size=14,
+        in_chans=3,
+        enc_embed_dim=1024,
+        apply_pe=False,
+    )
+
+    # Test dummy image input
+    dummy_image = torch.randn(1, 3, 518, 518)
+    patch_embedder_output = patch_embedder(ViTEncoderInput(data_norm_type="dense_rep_encoder", image=dummy_image))
+    assert patch_embedder_output.features.shape == (
+        1,
+        1024,
+        37,
+        37,
+    ), "Output features must have shape (1, 1024, 37, 37)"
+
+    # Init Dense Representation Encoder for non-image data as input
+    patch_embedder = DenseRepresentationEncoder(
+        name="dense_rep_encoder",
+        data_norm_type="dense_rep_encoder",
+        input_size_for_pe=518,
+        patch_size=14,
+        in_chans=6,
+        enc_embed_dim=1024,
+    )
+
+    # Init Dense Representation Encoder for single channel input
+    patch_embedder = DenseRepresentationEncoder(
+        name="dense_rep_encoder",
+        data_norm_type="dense_rep_encoder",
+        input_size_for_pe=518,
+        patch_size=14,
+        in_chans=1,
+        enc_embed_dim=1024,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        apply_pe=True,
+    )
+
+    # Test dummy non-image input
+    dummy_image = torch.randn(1, 1, 980, 980)
+    patch_embedder_output = patch_embedder(ViTEncoderNonImageInput(data=dummy_image))
+    assert patch_embedder_output.features.shape == (
+        1,
+        1024,
+        70,
+        70,
+    ), "Output features must have shape (1, 1024, 70, 70)"
+
+    print("All variants of Dense Representation Encoder have been initialized successfully!")

vismatch/third_party/UFM/UniCeption/uniception/models/encoders/dinov2.py

@@ -0,0 +1,333 @@
+"""
+Encoder Class for DINOv2
+"""
+
+from typing import List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from uniception.models.encoders.base import UniCeptionViTEncoderBase, ViTEncoderInput, ViTEncoderOutput
+from uniception.models.utils.intermediate_feature_return import IntermediateFeatureReturner
+
+
+class DINOv2Encoder(UniCeptionViTEncoderBase):
+    "UniCeption DINOv2 Encoder"
+
+    def __init__(
+        self,
+        name: str,
+        data_norm_type: str = "dinov2",
+        patch_size: int = 14,
+        size: str = "large",
+        with_registers: bool = False,
+        pretrained_checkpoint_path: str = None,
+        torch_hub_force_reload: bool = False,
+        gradient_checkpointing: bool = False,
+        keep_first_n_layers: Optional[int] = None,
+        use_pytorch_sdpa=True,
+        *args,
+        **kwargs,
+    ):
+        """
+        DINOv2 Encoder for extracting spatial features from images.
+
+        Args:
+            name (str): Name of the encoder.
+            data_norm_type (str): Image normalization type. Default: "dinov2"
+            patch_size (int): Patch size for the encoder. Default: 14
+            size (str): Size variant of the DINOv2 model. Options: ["small", "base", "large", "giant"]. Default: "large"
+            with_registers (bool): Whether to use the DINOv2 model with registers. Default: False
+            pretrained_checkpoint_path (str): Path to the pretrained checkpoint if using custom trained version of DINOv2. Default: None
+            torch_hub_force_reload (bool): Whether to force reload the model from torch hub. Default: False
+            gradient_checkpointing (bool): Whether to use gradient checkpointing to save GPU memory during backward call. Default: False
+            keep_first_n_layers (Optional[int]): If specified, only the first n layers of the model will be kept. Default: None
+            use_pytorch_sdpa (bool): Whether to use PyTorch native SDPA for attention layers. Default: True
+        """
+        # Init the base class
+        name = name if not with_registers else f"{name}_reg"
+        super().__init__(
+            name=name,
+            data_norm_type=data_norm_type,
+            patch_size=patch_size,
+            gradient_checkpointing=gradient_checkpointing,
+            *args,
+            **kwargs,
+        )
+
+        # Init the DINOv2 Encoder specific attributes
+        self.version = size
+        self.with_registers = with_registers
+        self.enc_embed_dim = {"small": 384, "base": 768, "large": 1024, "giant": 1536}[self.version]
+
+        # Define DINOv2 model factory
+        DINO_MODELS = {
+            # No registers
+            False: {
+                "small": "dinov2_vits14",
+                "base": "dinov2_vitb14",
+                "large": "dinov2_vitl14",
+                "giant": "dinov2_vitg14",
+            },
+            # With registers
+            True: {
+                "small": "dinov2_vits14_reg",
+                "base": "dinov2_vitb14_reg",
+                "large": "dinov2_vitl14_reg",
+                "giant": "dinov2_vitg14_reg",
+            },
+        }
+
+        # Load the pretrained DINOv2 model from torch hub
+        print(f"Loading pretrained {DINO_MODELS[self.with_registers][self.version]} from torch hub")
+        try:  # Requires internet access
+            self.model = torch.hub.load(
+                "facebookresearch/dinov2",
+                DINO_MODELS[self.with_registers][self.version],
+                force_reload=torch_hub_force_reload,
+            )
+        except:  # Load from cache
+            self.model = torch.hub.load("facebookresearch/dinov2", DINO_MODELS[self.with_registers][self.version])
+
+        del (
+            self.model.mask_token
+        )  # This parameter is unused in producing patch features, and will lead to unused parameters
+
+        # Keep only the first n layers of the model if keep_first_n_layers is specified
+        if keep_first_n_layers is not None:
+            self.model.blocks = nn.ModuleList(self.model.blocks[:keep_first_n_layers])
+
+        # Use Native Torch SDPA for attention layers if specified (instead of DINOv2's XFormers)
+        if use_pytorch_sdpa:
+            self.enable_pytorch_native_sdpa()
+
+        # Wrap the transformer blocks with support for gradient checkpointing if required
+        if self.gradient_checkpointing:
+            for i in range(len(self.model.blocks)):
+                self.model.blocks[i] = self.wrap_module_with_gradient_checkpointing(self.model.blocks[i])
+
+        # Load the custom pretrained checkpoint if provided
+        if pretrained_checkpoint_path:
+            print(f"Loading custom pretrained DINOv2 checkpoint from {pretrained_checkpoint_path}")
+            ckpt = torch.load(pretrained_checkpoint_path, weights_only=False)
+            print(self.load_state_dict(ckpt["model"]))
+
+    def enable_pytorch_native_sdpa(self):
+        "Enable PyTorch native SDPA for attention layers"
+        for i in range(len(self.model.blocks)):
+            self.model.blocks[i].attn = self.wrap_dinov2_attention_with_sdpa(self.model.blocks[i].attn)
+
+    def wrap_dinov2_attention_with_sdpa(self, module: nn.Module):
+        "Wrap DINOv2 attention module with PyTorch native SDPA"
+        assert torch.__version__ >= "2.0", "SDPA requires PyTorch 2.0 or later"
+
+        class _AttentionWrapper(module.__class__):
+            "SDPA Attention Wrapper Class"
+
+            def forward(self, x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                B, N, C = x.shape
+                qkv = (
+                    self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+                )  # (3, B, H, N, C // H)
+
+                q, k, v = torch.unbind(qkv, 0)  # (B, H, N, C // H)
+
+                x = F.scaled_dot_product_attention(q, k, v, attn_bias)
+                x = x.permute(0, 2, 1, 3).reshape(B, N, C)
+
+                x = self.proj(x)
+                x = self.proj_drop(x)
+                return x
+
+        module.__class__ = _AttentionWrapper
+        return module
+
+    def forward(self, encoder_input: ViTEncoderInput) -> ViTEncoderOutput:
+        """
+        DINOv2 Encoder Forward Pass
+
+        Args:
+            encoder_input (ViTEncoderInput): Input data for the encoder. Input data must contain image normalization type and normalized image tensor.
+
+        Returns:
+            ViTEncoderOutput: Output data from the encoder.
+        """
+        # Check image normalization type
+        self._check_data_normalization_type(encoder_input.data_norm_type)
+
+        # Check the dtype and shape of the input image
+        assert isinstance(encoder_input.image, torch.Tensor), "Input must be a torch.Tensor"
+        assert encoder_input.image.ndim == 4, "Input must be of shape (B, C, H, W)"
+        batch_size, channels, height, width = encoder_input.image.shape
+        assert channels == 3, "Input must have 3 channels"
+        assert (
+            height % self.patch_size == 0 and width % self.patch_size == 0
+        ), f"Input shape must be divisible by patch size: {self.patch_size}"
+
+        # Extract the features from the DINOv2 model
+        features = self.model.forward_features(encoder_input.image)["x_norm_patchtokens"]
+
+        # Resize the features to the expected shape
+        # (B x Num_patches x Embed_dim) -> (B x Embed_dim x H / Patch_Size x W / Patch_Size)
+        features = features.permute(0, 2, 1)
+        features = features.reshape(
+            -1, self.enc_embed_dim, height // self.patch_size, width // self.patch_size
+        ).contiguous()
+
+        return ViTEncoderOutput(features=features)
+
+
+class DINOv2IntermediateFeatureReturner(DINOv2Encoder, IntermediateFeatureReturner):
+    "Intermediate Feature Returner for UniCeption DINOv2 Encoder"
+
+    def __init__(
+        self,
+        name: str,
+        data_norm_type: str = "dinov2",
+        patch_size: int = 14,
+        size: str = "large",
+        with_registers: bool = False,
+        pretrained_checkpoint_path: str = None,
+        indices: Optional[Union[int, List[int]]] = 1,
+        keep_first_n_layers: Optional[int] = None,
+        norm_intermediate: bool = True,
+        *args,
+        **kwargs,
+    ):
+        """
+        DINOv2 Encoder for extracting spatial features from images.
+
+        Args:
+            name (str): Name of the encoder.
+            data_norm_type (str): Image normalization type. Default: "dinov2"
+            patch_size (int): Patch size for the encoder. Default: 14
+            size (str): Size variant of the DINOv2 model. Options: ["small", "base", "large", "giant"]
+            with_registers (bool): Whether to use the DINOv2 model with registers.
+            pretrained_checkpoint_path (str): Path to the pretrained checkpoint if using custom trained version of DINOv2.
+            indices (Optional[Union[int, List[int]]], optional): Indices of the layers to return. Defaults to 1. Options:
+            - int: Return the last n layers.
+            - List[int]: Return the intermediate layers at the specified indices.
+            keep_first_n_layers (Optional[int], optional): If specified, only the first n layers of the model will be kept. Defaults to None.
+            norm_intermediate (bool, optional): Whether to normalize the intermediate features. Defaults to True.
+        """
+        # Init the base classes
+        DINOv2Encoder.__init__(
+            self,
+            name=name,
+            data_norm_type=data_norm_type,
+            patch_size=patch_size,
+            size=size,
+            with_registers=with_registers,
+            keep_first_n_layers=keep_first_n_layers,
+            pretrained_checkpoint_path=pretrained_checkpoint_path,
+            *args,
+            **kwargs,
+        )
+        IntermediateFeatureReturner.__init__(
+            self,
+            indices=indices,
+            norm_intermediate=norm_intermediate,
+        )
+
+    def forward(self, encoder_input: ViTEncoderInput) -> List[ViTEncoderOutput]:
+        """
+        DINOv2 Encoder Forward Pass with Intermediate Feature Return
+
+        Args:
+            encoder_input (ViTEncoderInput): Input data for the encoder. Input data must contain image normalization type and normalized image tensor.
+
+        Returns:
+            List[ViTEncoderOutput]: Output data from the encoder. Returns a list of intermediate features.
+        """
+        # Check image normalization type
+        self._check_data_normalization_type(encoder_input.data_norm_type)
+
+        # Check the dtype and shape of the input image
+        assert isinstance(encoder_input.image, torch.Tensor), "Input must be a torch.Tensor"
+        assert encoder_input.image.ndim == 4, "Input must be of shape (B, C, H, W)"
+        batch_size, channels, height, width = encoder_input.image.shape
+        assert channels == 3, "Input must have 3 channels"
+        assert (
+            height % self.patch_size == 0 and width % self.patch_size == 0
+        ), f"Input shape must be divisible by patch size: {self.patch_size}"
+
+        if self.indices is None:
+            self.indices = range(len(self.model.blocks))
+
+        # Extract the intermediate features from the DINOv2 model
+        intermediate_features = self.model.get_intermediate_layers(
+            encoder_input.image, n=self.indices, reshape=True, norm=self.norm_intermediate
+        )
+
+        # Convert the intermediate features to a list of ViTEncoderOutput
+        intermediate_features = [ViTEncoderOutput(features=features) for features in intermediate_features]
+
+        return intermediate_features
+
+
+if __name__ == "__main__":
+    # Init different variants of DINOv2
+    for size in ["small", "base", "large", "giant"]:
+        for with_registers in [False, True]:
+            name = f"dinov2_{size}"
+            dinov2_encoder = DINOv2Encoder(name=name, size=size, with_registers=with_registers)
+
+    # Init the custom pretrained DINOv2 encoders
+    for size in ["small", "base", "large"]:
+        pretrained_checkpoints_dict = {
+            "small": "../../../checkpoints/encoders/DINOv2_ViTS_DepthAnythingV2.pth",
+            "base": "../../../checkpoints/encoders/DINOv2_ViTB_DepthAnythingV2.pth",
+            "large": "../../../checkpoints/encoders/DINOv2_ViTL_DepthAnythingV2.pth",
+        }
+        name = f"dinov2_dav2_{size}"
+        dinov2_encoder = DINOv2Encoder(
+            name=name, size=size, with_registers=False, pretrained_checkpoint_path=pretrained_checkpoints_dict[size]
+        )
+
+    print("All DINOv2 Encoders have been initialized successfully!")
+
+    # Intermediate Feature Returner Tests
+    print("Running Intermediate Feature Returner Tests...")
+
+    # Run the intermediate feature returner with last-n index
+    dinov2_intermediate_feature_returner = DINOv2IntermediateFeatureReturner(
+        name="dinov2_base", size="base", indices=6
+    )  # Last 6 layers
+    dummy_input = ViTEncoderInput(image=torch.randn(1, 3, 224, 224), data_norm_type="dinov2")
+    output = dinov2_intermediate_feature_returner(dummy_input)
+    assert isinstance(output, list), "Output must be a list of intermediate features"
+    assert isinstance(output[0], ViTEncoderOutput), "Output must be a list of ViTEncoderOutput"
+    assert len(output) == 6, "Output must have length of intermediate features equal to the number of indices"
+
+    # Run the intermediate feature returner with specific indices
+    dinov2_intermediate_feature_returner = DINOv2IntermediateFeatureReturner(
+        name="dinov2_base", size="base", indices=[0, 2, 4, 6]
+    )  # Specific layers
+    dummy_input = ViTEncoderInput(image=torch.randn(1, 3, 224, 224), data_norm_type="dinov2")
+    output = dinov2_intermediate_feature_returner(dummy_input)
+    assert isinstance(output, list), "Output must be a list of intermediate features"
+    assert isinstance(output[0], ViTEncoderOutput), "Output must be a list of ViTEncoderOutput"
+    assert len(output) == 4, "Output must have length of intermediate features equal to the number of indices"
+
+    print("All Intermediate Feature Returner Tests have passed successfully!")
+
+    from uniception.models.encoders.utils import profile_encoder
+
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    # Profile the DINOv2 Encoder
+    dinov2_encoder = DINOv2Encoder(
+        name="dinov2_large", size="large", use_pytorch_sdpa=True, gradient_checkpointing=True, keep_first_n_layers=12
+    ).cuda()
+    dummy_input = ViTEncoderInput(image=torch.randn(24, 3, 560, 420).cuda(), data_norm_type="dinov2")
+
+    class Profiler:
+        @profile_encoder(num_warmup=3, num_runs=20, autocast_precision="bfloat16", use_compile=True, dynamic=False)
+        def run_fn(self):
+            output = dinov2_encoder(dummy_input)
+            return output
+
+    profiler = Profiler()
+    profiler.run_fn()