init

InfiniDepth/__init__.py

@@ -0,0 +1 @@
+"""InfiniDepth package."""

InfiniDepth/model/__init__.py

@@ -0,0 +1,9 @@
+from .registry import MODEL_REGISTRY, register_model
+from .model import InfiniDepth, InfiniDepth_DC
+
+__all__ = [
+    "MODEL_REGISTRY",
+    "register_model",
+    "InfiniDepth",
+    "InfiniDepth_DC",
+]

InfiniDepth/model/block/__init__.py

@@ -0,0 +1 @@
+"""Core building blocks for InfiniDepth models."""

InfiniDepth/model/block/common.py

@@ -0,0 +1,43 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+
+from typing import Type
+
+
+class MLPBlock(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        mlp_dim: int,
+        act: Type[nn.Module] = nn.GELU,
+    ) -> None:
+        super().__init__()
+        self.lin1 = nn.Linear(embedding_dim, mlp_dim)
+        self.lin2 = nn.Linear(mlp_dim, embedding_dim)
+        self.act = act()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lin2(self.act(self.lin1(x)))
+
+
+# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
+# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119  # noqa
+class LayerNorm2d(nn.Module):
+    def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(num_channels))
+        self.bias = nn.Parameter(torch.zeros(num_channels))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        x = self.weight[:, None, None] * x + self.bias[:, None, None]
+        return x

InfiniDepth/model/block/config.py

@@ -0,0 +1,5 @@
+dinov3_model_configs = {
+    "vitl16":{'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024], 'layer_idxs': [4, 11, 17, 23]},
+    "vith16plus": {'encoder': 'vith', 'features': 384, 'out_channels': [1280, 1280, 1280, 1280], 'layer_idxs': [7, 15, 23, 31]},
+}
+

InfiniDepth/model/block/convolution.py

@@ -0,0 +1,229 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+from typing import Callable
+import collections
+from torch import Tensor
+from itertools import repeat
+
+
+def bilinear_sampler(input, coords, align_corners=True, padding_mode="border"):
+    r"""Sample a tensor using bilinear interpolation
+
+    `bilinear_sampler(input, coords)` samples a tensor :attr:`input` at
+    coordinates :attr:`coords` using bilinear interpolation. It is the same
+    as `torch.nn.functional.grid_sample()` but with a different coordinate
+    convention.
+
+    The input tensor is assumed to be of shape :math:`(B, C, H, W)`, where
+    :math:`B` is the batch size, :math:`C` is the number of channels,
+    :math:`H` is the height of the image, and :math:`W` is the width of the
+    image. The tensor :attr:`coords` of shape :math:`(B, H_o, W_o, 2)` is
+    interpreted as an array of 2D point coordinates :math:`(x_i,y_i)`.
+
+    Alternatively, the input tensor can be of size :math:`(B, C, T, H, W)`,
+    in which case sample points are triplets :math:`(t_i,x_i,y_i)`. Note
+    that in this case the order of the components is slightly different
+    from `grid_sample()`, which would expect :math:`(x_i,y_i,t_i)`.
+
+    If `align_corners` is `True`, the coordinate :math:`x` is assumed to be
+    in the range :math:`[0,W-1]`, with 0 corresponding to the center of the
+    left-most image pixel :math:`W-1` to the center of the right-most
+    pixel.
+
+    If `align_corners` is `False`, the coordinate :math:`x` is assumed to
+    be in the range :math:`[0,W]`, with 0 corresponding to the left edge of
+    the left-most pixel :math:`W` to the right edge of the right-most
+    pixel.
+
+    Similar conventions apply to the :math:`y` for the range
+    :math:`[0,H-1]` and :math:`[0,H]` and to :math:`t` for the range
+    :math:`[0,T-1]` and :math:`[0,T]`.
+
+    Args:
+        input (Tensor): batch of input images.
+        coords (Tensor): batch of coordinates.
+        align_corners (bool, optional): Coordinate convention. Defaults to `True`.
+        padding_mode (str, optional): Padding mode. Defaults to `"border"`.
+
+    Returns:
+        Tensor: sampled points.
+    """
+
+    sizes = input.shape[2:]
+
+    assert len(sizes) in [2, 3]
+
+    if len(sizes) == 3:
+        # t x y -> x y t to match dimensions T H W in grid_sample
+        coords = coords[..., [1, 2, 0]]
+
+    if align_corners:
+        coords = coords * torch.tensor(
+            [2 / max(size - 1, 1) for size in reversed(sizes)], device=coords.device
+        )
+    else:
+        coords = coords * torch.tensor(
+            [2 / size for size in reversed(sizes)], device=coords.device
+        )
+
+    coords -= 1
+
+    return F.grid_sample(
+        input, coords, align_corners=align_corners, padding_mode=padding_mode
+    )
+
+
+def round_to_multiple_of_4(n):
+    return round(n / 4) * 4
+
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn="group", stride=1):
+        super(ResidualBlock, self).__init__()
+
+        self.conv1 = nn.Conv2d(
+            in_planes,
+            planes,
+            kernel_size=3,
+            padding=1,
+            stride=stride,
+            padding_mode="zeros",
+        )
+        self.conv2 = nn.Conv2d(
+            planes, planes, kernel_size=3, padding=1, padding_mode="zeros"
+        )
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == "group":
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+
+        elif norm_fn == "batch":
+            self.norm1 = nn.BatchNorm2d(planes)
+            self.norm2 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.BatchNorm2d(planes)
+
+        elif norm_fn == "instance":
+            self.norm1 = nn.InstanceNorm2d(planes)
+            self.norm2 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == "none":
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            if not stride == 1:
+                self.norm3 = nn.Sequential()
+
+        if stride == 1:
+            self.downsample = None
+
+        else:
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3
+            )
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x + y)
+
+        
+class BasicEncoder(nn.Module):
+    def __init__(self, input_dim=3, output_dim=128, stride=4):
+        super(BasicEncoder, self).__init__()
+        self.stride = stride
+        self.norm_fn = "instance"
+        self.in_planes = output_dim // 2
+        self.norm1 = nn.InstanceNorm2d(self.in_planes)
+        self.norm2 = nn.InstanceNorm2d(output_dim * 2)
+
+        self.conv1 = nn.Conv2d(
+            input_dim,
+            self.in_planes,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            padding_mode="zeros",
+        )
+        self.relu1 = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(output_dim // 2, stride=1)
+        self.layer2 = self._make_layer(output_dim // 4 * 3, stride=2)
+        self.layer3 = self._make_layer(output_dim, stride=2)
+        self.layer4 = self._make_layer(output_dim, stride=2)
+
+        self.conv2 = nn.Conv2d(
+            output_dim * 3 + output_dim // 4,
+            output_dim * 2,
+            kernel_size=3,
+            padding=1,
+            padding_mode="zeros",
+        )
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(output_dim * 2, output_dim, kernel_size=1)
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.InstanceNorm2d)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        _, _, H, W = x.shape
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        a = self.layer1(x)
+        b = self.layer2(a)
+        c = self.layer3(b)
+        d = self.layer4(c)
+
+        def _bilinear_intepolate(x):
+            return F.interpolate(
+                x,
+                (H // self.stride, W // self.stride),
+                mode="bilinear",
+                align_corners=True,
+            )
+
+        a = _bilinear_intepolate(a)
+        b = _bilinear_intepolate(b)
+        c = _bilinear_intepolate(c)
+        d = _bilinear_intepolate(d)
+
+        x = self.conv2(torch.cat([a, b, c, d], dim=1))
+        x = self.norm2(x)
+        x = self.relu2(x)
+        x = self.conv3(x)
+        return x

InfiniDepth/model/block/implicit_decoder.py

@@ -0,0 +1,179 @@
+import enum
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import sys
+from grpc import insecure_channel
+from sympy import use
+from pathlib import Path
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+class MLP(nn.Module):
+    def __init__(self, in_dim, out_dim, hidden_list, output_act='elu'):
+        super().__init__()
+        layers = []
+        lastv = in_dim
+        for hidden in hidden_list:
+            layers += [nn.Linear(lastv, hidden), nn.ReLU()]
+            lastv = hidden
+
+        if out_dim is not None:
+            layers.append(nn.Linear(lastv, out_dim))
+            act = {
+                "sigmoid": nn.Sigmoid(),
+                "relu": nn.ReLU(),
+                "elu": nn.ELU(),
+            }.get(output_act, nn.Identity())
+            layers.append(act)
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+    
+
+class ImplicitHead(nn.Module):
+    """
+    Implicit head that fuses DINOv3 semantic features and BasicEncoder low-level features.
+
+    Args:
+        hidden_dim: DINOv3 feature dimension (e.g., 1024)
+        basic_dim: BasicEncoder feature dimension (e.g., 128)
+        fusion_type: Feature fusion strategy
+            - "concat": Simple concatenation
+            - "cross_attn": Cross-attention between features
+            - "gated": Gated fusion with learnable weights
+        out_dim: Output dimension (1 for depth)
+        hidden_list: MLP hidden layer dimensions
+    """
+    def __init__(
+            self,
+            hidden_dim,  # 1024 for DINOv3
+            basic_dim=128,  # BasicEncoder output dim
+            fusion_type="gated",  # concat, gated
+            out_dim=1,
+            hidden_list=[1024, 256, 32],
+            ):
+
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.basic_dim = basic_dim
+        self.fusion_type = fusion_type
+
+        # Determine input dimension based on fusion type
+        if fusion_type == "concat":
+            # Simple concatenation
+            in_channels = hidden_dim + basic_dim
+        elif fusion_type == "gated":
+            # Gated fusion with learnable weights
+            self.gate_proj = nn.Linear(basic_dim, hidden_dim)
+            self.gate = nn.Sequential(
+                nn.Linear(hidden_dim * 2, hidden_dim),
+                nn.Sigmoid()
+            )
+            in_channels = hidden_dim
+        else:
+            raise ValueError(f"Unknown fusion_type: {fusion_type}")
+
+        self.out_layer = MLP(
+            in_dim=in_channels,
+            out_dim=out_dim,
+            hidden_list=hidden_list,
+            output_act='elu'
+        )
+
+    def _encode_feat(self, features, patch_h, patch_w):
+        """Extract DINOv3 feature map."""
+        x = features[-1][0]
+        out_feat = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
+        return out_feat
+
+    def _decode_dpt(self, feat, basic_feat, coord):
+        """
+        Query features at given coordinates and fuse them.
+
+        Args:
+            feat: DINOv3 feature map [B, hidden_dim, H_dino, W_dino]
+            basic_feat: BasicEncoder feature map [B, basic_dim, H_basic, W_basic]
+            coord: Query coordinates [B, N, 2] in range [-1, 1]
+
+        Returns:
+            pred: Predicted depth [B, N, 1]
+        """
+        coord_ = coord.clone()
+        coord_.clamp_(-1 + 1e-6, 1 - 1e-6)
+
+        # Sample DINOv3 features at query coordinates
+        q_feat_dino = F.grid_sample(
+            feat, coord_.flip(-1).unsqueeze(1),
+            mode='bilinear', align_corners=False
+        )[:, :, 0, :].permute(0, 2, 1)  # [B, N, hidden_dim]
+
+        # Sample BasicEncoder features at query coordinates (if available)
+        if basic_feat is not None:
+            q_feat_basic = F.grid_sample(
+                basic_feat, coord_.flip(-1).unsqueeze(1),
+                mode='bilinear', align_corners=False
+            )[:, :, 0, :].permute(0, 2, 1)  # [B, N, basic_dim]
+
+            # Fuse features based on fusion type
+            q_feat_fused = self._fuse_features(q_feat_dino, q_feat_basic)
+        else:
+            # If no basic features, use only DINOv3
+            q_feat_fused = q_feat_dino
+
+        # Predict depth
+        pred = self.out_layer(q_feat_fused)
+        return pred
+
+    def _fuse_features(self, feat_dino, feat_basic):
+        """
+        Fuse DINOv3 and BasicEncoder features.
+
+        Args:
+            feat_dino: [B, N, hidden_dim]
+            feat_basic: [B, N, basic_dim]
+
+        Returns:
+            fused_feat: [B, N, fused_dim]
+        """
+        if self.fusion_type == "concat":
+            # Simple concatenation
+            return torch.cat([feat_dino, feat_basic], dim=-1)
+
+        elif self.fusion_type == "gated":
+            # Gated fusion with learnable weights
+            feat_basic_proj = self.gate_proj(feat_basic)  # [B, N, hidden_dim]
+            gate_input = torch.cat([feat_dino, feat_basic_proj], dim=-1)
+            gate_weights = self.gate(gate_input)  # [B, N, hidden_dim]
+            return gate_weights * feat_dino + (1 - gate_weights) * feat_basic_proj
+
+    def forward(self, features, basic_feat, patch_h, patch_w, coords):
+        """
+        Forward pass.
+
+        Args:
+            features: DINOv3 features from backbone
+            basic_feat: BasicEncoder features [B, basic_dim, H/4, W/4]
+            patch_h, patch_w: DINOv3 feature map spatial size
+            coords: Query coordinates [B, N, 2]
+
+        Returns:
+            dpt_pred: Predicted depth [B, N, 1]
+        """
+        # Extract DINOv3 feature map
+        feat = self._encode_feat(features, patch_h, patch_w)  # [B, hidden_dim, H/14, W/14]
+
+        # Query and fuse features at coordinates
+        dpt_pred = self._decode_dpt(feat, basic_feat, coords)
+
+        return dpt_pred

InfiniDepth/model/block/pe.py

@@ -0,0 +1,222 @@
+import torch
+import torch.nn as nn
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Any, Optional, Tuple, Dict
+
+if torch.cuda.is_available():
+    acc_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
+else:
+    acc_dtype = torch.float16
+
+POS_EMB_REGISTRY = {}
+
+def register_pos_emb(name):
+    def decorator(cls):
+        POS_EMB_REGISTRY[name.lower()] = cls
+        return cls
+    return decorator
+
+        
+@register_pos_emb("dct")
+class DctPositionEmbedding(nn.Module):
+    """
+    Only supports 2D separable DCT encoding for query coordinates coords: [B, N, 2]:
+      Phi(x,y)[fx,fy] = cos(pi * fx * x) * cos(pi * fy * y) * 1/(1+fx*fy)
+    Convention: coords should be in the range [0, 1].
+    """
+    def __init__(self, max_freqs: int = 8):
+        super().__init__()
+        self.max_freqs = max_freqs
+
+        freqs = torch.arange(max_freqs).float()              # [F] -> 0..F-1
+        fx = freqs.view(-1, 1)                               # [F,1]
+        fy = freqs.view(1, -1)                               # [1,F]
+        coeffs = (1.0 + fx * fy) ** -1                       # [F,F]
+
+        self.register_buffer("_freqs_1d", freqs, persistent=False)
+        self.register_buffer("_coeffs_2d", coeffs, persistent=False)
+
+    def forward(self, coords: torch.Tensor) -> torch.Tensor:
+        """
+        coords: [B, N, 2], value range should be [0,1]
+        return: [B, N, F^2]
+        """
+        assert coords.dim() == 3 and coords.size(-1) == 2, "coords must be [B, N, 2]"
+        B, N, _ = coords.shape
+        device, dtype = coords.device, coords.dtype
+
+        freqs = self._freqs_1d.to(device=device, dtype=dtype)    # [F]
+        coeffs = self._coeffs_2d.to(device=device, dtype=dtype)  # [F,F]
+        F = freqs.numel() # frequency dimension = max_freqs
+
+        x = coords[..., 0:1]                                     # [B,N,1]
+        y = coords[..., 1:2]                                     # [B,N,1]
+        dct_x = torch.cos(math.pi * x * freqs.view(1, 1, F))     # [B,N,F]
+        dct_y = torch.cos(math.pi * y * freqs.view(1, 1, F))     # [B,N,F]
+
+        out = dct_x.unsqueeze(-1) * dct_y.unsqueeze(-2)          # [B,N,F,F]
+        out = out * coeffs.view(1, 1, F, F)                      # [B,N,F,F]
+        dct_emb = out.reshape(B, N, F * F)                       # [B,N,F^2]
+
+        return dct_emb                                           # [B,N,F^2]
+
+
+@register_pos_emb("random")
+class RandomPositionEmbedding(nn.Module):
+    """
+    Positional encoding using random spatial frequencies.
+    """
+
+    patch_size = 14
+
+    def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None, image_pe_method: str = "patch") -> None:
+        super().__init__()
+        if scale is None or scale <= 0.0:
+            scale = 1.0
+        self.register_buffer(
+            "positional_encoding_gaussian_matrix",
+            scale * torch.randn((2, num_pos_feats)),
+        )
+
+        self.image_pe_method = image_pe_method
+        if self.image_pe_method == "image":
+            # self.patch_embed = nn.Conv2d(num_pos_feats*2, num_pos_feats*2, kernel_size=self.patch_size, stride=self.patch_size)
+            self.patch_embed = nn.Sequential(
+                nn.Conv2d(num_pos_feats * 2, num_pos_feats // 2, kernel_size=2, stride=2),
+                nn.ReLU(),
+                nn.Conv2d(
+                    num_pos_feats // 2,
+                    num_pos_feats * 2,
+                    kernel_size=self.patch_size // 2,
+                    stride=self.patch_size // 2,
+                ),
+            )
+
+    def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
+        """Positionally encode points that are normalized to [0,1]."""
+        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
+        coords = 2 * coords - 1  # [0,1] --> [-1,1], equivalent to align_corners=False after this transform
+        coords = coords @ self.positional_encoding_gaussian_matrix
+        coords = 2 * np.pi * coords
+        # outputs d_1 x ... x d_n x C shape
+        return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
+
+    def forward_encoding(self, size: Tuple[int, int]) -> torch.Tensor:
+        """Generate positional encoding for a grid of the specified size."""
+        h, w = size
+        device: Any = self.positional_encoding_gaussian_matrix.device
+        grid = torch.ones((h, w), device=device, dtype=torch.float32)
+        y_embed = grid.cumsum(dim=0) - 0.5
+        x_embed = grid.cumsum(dim=1) - 0.5
+        y_embed = y_embed / h
+        x_embed = x_embed / w
+        pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))  # HxWx2 -> HxWxC
+        return pe.permute(2, 0, 1)  # C x H x W
+
+    def forward(self, size: Tuple[int, int]) -> torch.Tensor:
+        if self.image_pe_method == "patch":
+            return self.forward_encoding(size)
+        elif self.image_pe_method == "image":
+            pe_encoding = self.forward_encoding(size)
+            pe_encoding_high = self.forward_encoding((size[0] * self.patch_size, size[1] * self.patch_size))
+            return pe_encoding + self.patch_embed(pe_encoding_high[None])[0]
+
+    def forward_with_coords(self, coords_input: torch.Tensor, image_size: Tuple[int, int]) -> torch.Tensor:
+        """Positionally encode points that are not normalized to [0,1]."""
+        coords = coords_input.clone()
+        coords[:, :, 0] = coords[:, :, 0] / image_size[1]
+        coords[:, :, 1] = coords[:, :, 1] / image_size[0]
+        return self._pe_encoding(coords.to(torch.float))  # B x N x C
+    
+
+@register_pos_emb("rope")
+class RoPEPositionEmbedding(nn.Module):
+    """2D Rotary Position Embedding with support for continuous coordinates.
+
+    For each coordinate p (can be float), directly compute θ = p * inv_freq, then derive cos/sin.
+    """
+    def __init__(self, frequency: float = 100.0, scaling_factor: float = 1.0):
+        super().__init__()
+        self.base_frequency = frequency
+        self.scaling_factor = scaling_factor
+        # Cache the inv_freq vector: key = feature_dim
+        self._inv_freq_cache: Dict[int, torch.Tensor] = {}
+
+    def _get_inv_freq(self, dim: int, device: torch.device, dtype: torch.dtype):
+        """
+        Computes frequency components for rotary embeddings.
+        Returns an inv_freq vector of length dim/2, in the form 1 / base_freq^(2i/d)
+        """
+        if dim not in self._inv_freq_cache:
+            # Use frequencies on even dimensions only
+            exponents = torch.arange(0, dim, 2, device=device).float() / dim
+            inv_freq = 1.0 / (self.base_frequency ** exponents)
+            self._inv_freq_cache[dim] = inv_freq.to(dtype)
+        return self._inv_freq_cache[dim]
+
+    @staticmethod
+    def _rotate_features(x: torch.Tensor) -> torch.Tensor:
+        """Rotation: split [u0, u1, u2, u3,...] into two halves and concatenate (-v, u)."""
+        D = x.shape[-1]
+        x1, x2 = x[..., : D//2], x[..., D//2 :]
+        return torch.cat((-x2, x1), dim=-1)
+
+    def _apply_1d_rope_continuous(
+        self, 
+        x: torch.Tensor,          # [B, N, d_half]
+        pos: torch.Tensor,        # [B, N] floating-point coordinates
+        inv_freq: torch.Tensor    # [d_half]
+    ) -> torch.Tensor:
+        # 1) Compute angles: [B, N, d_half] = outer(pos, inv_freq)
+        #    pos.unsqueeze(-1): [B, N, 1], inv_freq.unsqueeze(0): [1, d_half]
+        angles = pos.unsqueeze(-1) * inv_freq.unsqueeze(0)
+        # 2) Duplicate to double dimension: [B, N, d_half*2]
+        angles = torch.cat([angles, angles], dim=-1)
+
+        # 3) Compute cos/sin and expand to [B, N, D]
+        cos = angles.cos()
+        sin = angles.sin()
+
+        # 4) Apply rotation
+        return x * cos + self._rotate_features(x) * sin
+
+    def forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        """
+        tokens:    [B, N, dim]   
+        positions: [B, N, 2]    continuous coords: (y,x)
+        """
+        B, N, D = tokens.shape
+        assert D % 2 == 0, "Feature dimension must be even"
+
+        assert positions.shape == (B, N, 2), "positions must be [B, N, 2]"
+
+        # Allocate half of the features to each direction
+        d_half = D // 2
+
+        # Get the inv_freq vector
+        inv_freq = self._get_inv_freq(d_half, tokens.device, tokens.dtype)  # [d_half]
+        # Split feature dimension into first and second halves
+        tok_v, tok_h = tokens[..., :d_half], tokens[..., d_half:]
+
+        # Apply RoPE separately on y and x directions, positions[0]--> y, positions[1]--> x
+        out_v = self._apply_1d_rope_continuous(tok_v, positions[..., 0], inv_freq)
+        out_h = self._apply_1d_rope_continuous(tok_h, positions[..., 1], inv_freq)
+
+        return torch.cat([out_v, out_h], dim=-1)
+
+
+def build_pos_emb(pos_emb_type="nerf", **kwargs):
+    pos_emb_type = pos_emb_type.lower()
+    if pos_emb_type not in POS_EMB_REGISTRY:
+        raise ValueError(f"Unknown pos_emb_type: {pos_emb_type}")
+    return POS_EMB_REGISTRY[pos_emb_type](**kwargs)
+
+
+
+
+
+

InfiniDepth/model/block/perceive_io.py

@@ -0,0 +1,274 @@
+from typing import Callable, Optional, Union
+import torch
+from einops import rearrange
+from ...utils.logger import Log
+from torch import Tensor, nn
+
+try:
+    from xformers.ops import memory_efficient_attention, unbind
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    Log.warning("xFormers not available")
+    XFORMERS_AVAILABLE = False
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        context_dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        pe: str = "normal",
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim, bias=qkv_bias)
+        self.qkv_context = nn.Linear(context_dim, context_dim * 2, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.pe = pe
+        if self.pe == "qk":
+            self.norm1 = nn.LayerNorm(dim // num_heads)
+            self.norm2 = nn.LayerNorm(dim // num_heads)
+
+    def forward(self, x: Tensor, context: Tensor) -> Tensor:
+        # x is the query tensor, context is the key/value tensor
+        B, N, C = x.shape
+        _, M, _ = context.shape
+
+        qkv_x = self.qkv(x).reshape(B, N, 1, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q_x = qkv_x[0] * self.scale
+
+        qkv_context = (
+            self.qkv_context(context).reshape(B, M, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        )
+        k_context, v_context = qkv_context[0], qkv_context[0]
+
+        # Cross-attention: query from x and key/value from context
+        attn = q_x @ k_context.transpose(-2, -1)
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v_context).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffCrossAttention(CrossAttention):
+    def forward(
+        self, x: Tensor, context: Tensor, x_pe: Tensor = None, context_pe: Tensor = None, attn_bias=None
+    ) -> Tensor:
+        if not XFORMERS_AVAILABLE:
+            assert attn_bias is None, "xFormers is required for nested tensors usage"
+            return super().forward(x, context)
+
+        B, N, C = x.shape
+        _, M, C_context = context.shape
+
+        qkv_x = self.qkv(x).reshape(B, N, 1, self.num_heads, C // self.num_heads)
+        (q_x,) = unbind(qkv_x, 2)
+
+        qkv_context = self.qkv_context(context).reshape(B, M, 2, self.num_heads, C_context // self.num_heads)
+        k_context, v_context = unbind(qkv_context, 2)
+
+        if self.pe == "qk":
+            q_x = self.norm1(q_x + rearrange(x_pe, "b n (m c) -> b n m c", m=self.num_heads))
+            k_context = self.norm2(k_context + rearrange(context_pe, "b n (m c) -> b n m c", m=self.num_heads))
+        elif self.pe == "apply":
+            pass
+        # Memory-efficient cross-attention
+        x = memory_efficient_attention(
+            q_x.to(dtype=v_context.dtype), k_context.to(dtype=v_context.dtype), v_context, attn_bias=attn_bias
+        )
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+# class Attention(nn.Module):
+#     def __init__(
+#         self,
+#         dim: int,
+#         num_heads: int = 8,
+#         qkv_bias: bool = False,
+#         proj_bias: bool = True,
+#         attn_drop: float = 0.0,
+#         proj_drop: float = 0.0,
+#     ) -> None:
+#         super().__init__()
+#         self.num_heads = num_heads
+#         head_dim = dim // num_heads
+#         self.scale = head_dim**-0.5
+
+#         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+#         self.attn_drop = nn.Dropout(attn_drop)
+#         self.proj = nn.Linear(dim, dim, bias=proj_bias)
+#         self.proj_drop = nn.Dropout(proj_drop)
+
+#     def forward(self, x: Tensor) -> 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)
+
+#         q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+#         attn = q @ k.transpose(-2, -1)
+
+#         attn = attn.softmax(dim=-1)
+#         attn = self.attn_drop(attn)
+
+#         x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+#         x = self.proj(x)
+#         x = self.proj_drop(x)
+#         return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        pe: str = "normal",
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.pe = pe
+        if self.pe == "qk":
+            self.norm1 = nn.LayerNorm(dim // num_heads)
+            self.norm2 = nn.LayerNorm(dim // num_heads)
+
+    def forward(self, x: Tensor, x_pe: Tensor = None) -> 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)
+
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+        if self.pe == "qk":
+            q = self.norm1(q + x_pe)
+            k = self.norm2(k + x_pe)
+        elif self.pe == "apply":
+            pass
+        attn = q @ k.transpose(-2, -1)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(Attention):
+    def forward(
+        self,
+        x: Tensor,
+        x_pe: Tensor = None,
+        attn_bias=None,
+    ) -> Tensor:
+        if not XFORMERS_AVAILABLE:
+            assert attn_bias is None, "xFormers is required for nested tensors usage"
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        q, k, v = unbind(qkv, 2)
+        if self.pe == "qk":
+            q = self.norm1(q + rearrange(x_pe, "b n (m c) -> b n m c", m=self.num_heads))
+            k = self.norm2(k + rearrange(x_pe, "b n (m c) -> b n m c", m=self.num_heads))
+        elif self.pe == "apply":
+            pass
+        # this is important
+        # as q, k after norm1/norm2 have different dtype
+        # which will cause error in memory_efficient_attention
+        x = memory_efficient_attention(q.to(dtype=v.dtype), k.to(dtype=v.dtype), v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)

InfiniDepth/model/block/prompt_models/__init__.py

@@ -0,0 +1,31 @@
+import torch.nn as nn
+from .sam import SAMPromptModel
+from .selfattn import SelfAttnPromptModel
+
+__all__ = [
+    "GeneralPromptModel",
+    "SelfAttnPromptModel",
+    "SAMPromptModel",
+]
+
+
+class GeneralPromptModel(nn.Module):
+    def __init__(self, prompt_stage=[3], **kwargs):
+        super().__init__()
+        self.prompt_stage = prompt_stage
+        self.prompt_idmap = {i: idx for idx, i in enumerate(self.prompt_stage)}
+        block = kwargs.get("block")
+        self.prompt_model = nn.ModuleList([block for _ in range(len(self.prompt_stage))])
+
+    def forward(self, features, prompt_depth, prompt_mask, patch_h, patch_w):
+        for i in range(len(features)):
+            if i not in self.prompt_stage: # prompt_stage = [3]
+                continue
+            features[i][0] = self.prompt_model[self.prompt_idmap[i]](
+                features[i][0],
+                prompt_depth,
+                prompt_mask,
+                patch_h,
+                patch_w,
+            )
+        return features

InfiniDepth/model/block/prompt_models/crossattn.py

@@ -0,0 +1,164 @@
+import torch
+import torch.nn as nn
+from typing import Tuple
+from ..perceive_io import LayerScale, MemEffCrossAttention, Mlp
+from .utils.pe_utils import PositionEmbeddingRandom
+from torch import Tensor
+
+
+class CrossAttnPromptModel(nn.Module):
+    def __init__(
+        self,
+        transformer_dim: int = 1024,
+        num_blocks: int = 1,
+        num_heads: int = 4,
+        pe: str = "normal",
+        image_pe_method: str = "patch",  # image
+        **kwargs,
+    ) -> None:
+        """
+        Predicts masks given an image and prompt embeddings, using a
+        transformer architecture.
+
+        Arguments:
+          transformer_dim (int): the channel dimension of the transformer
+          transformer (nn.Module): the transformer used to predict masks
+          num_multimask_outputs (int): the number of masks to predict
+            when disambiguating masks
+          activation (nn.Module): the type of activation to use when
+            upscaling masks
+          iou_head_depth (int): the depth of the MLP used to predict
+            mask quality
+          iou_head_hidden_dim (int): the hidden dimension of the MLP
+            used to predict mask quality
+        """
+        super().__init__()
+        self.pe = pe
+        pe_dim = transformer_dim // 2
+        if self.pe == "apply":
+            pe_dim = pe_dim // num_heads
+        self.pe_layer = PositionEmbeddingRandom(pe_dim, image_pe_method=image_pe_method)
+        self.prompt_blocks = nn.ModuleList(
+            [
+                CrossAttenPromptBlock(dim=transformer_dim, num_heads=num_heads, first_block=(i == 0), pe=pe)
+                for i in range(num_blocks)
+            ]
+        )
+        self.depth2feature = nn.Sequential(
+            nn.Linear(1, transformer_dim // 2),
+            nn.GELU(),
+            nn.Linear(transformer_dim // 2, transformer_dim),
+        )
+
+    def forward(
+        self,
+        image_embeddings: torch.Tensor,
+        prompt_depth: torch.Tensor,
+        prompt_mask: torch.Tensor,
+        patch_h: int,
+        patch_w: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Predict masks given image and prompt embeddings.
+        Arguments:
+          image_embeddings (torch.Tensor): the embeddings from the image encoder
+          image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
+          sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
+          dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
+          multimask_output (bool): Whether to return multiple masks or a single
+            mask.
+        Returns:
+          torch.Tensor: batched predicted masks
+          torch.Tensor: batched predictions of mask quality
+        """
+        B, _, H, W = prompt_depth.shape
+        image_pe = self.pe_layer((patch_h, patch_w)).permute(1, 2, 0)  # CxHxW -> HxWxC
+        prompt_embeddings_list = []
+        image_embeddings_list = []
+        for b in range(B):
+            valid_pts_num = (prompt_mask[b, 0] > 0.0).sum()
+            if valid_pts_num == 0:
+                image_embeddings_item = image_embeddings[b : (b + 1)]
+                image_embeddings_list.append(image_embeddings_item)
+                continue
+            sparse_depth_pos = (prompt_mask[b, 0] > 0.0).nonzero().float()
+            sparse_depth_pos[:, 0] = (sparse_depth_pos[:, 0] + 0.5) / H
+            sparse_depth_pos[:, 1] = (sparse_depth_pos[:, 1] + 0.5) / W
+            sparse_depth = prompt_depth[b, 0][prompt_mask[b, 0] > 0.0]
+            prompt_embeddings = self.depth2feature(sparse_depth[:, None])[None, ...]  # 1, N, C
+            prompt_pe = self.pe_layer._pe_encoding(sparse_depth_pos[None, :, [1, 0]])  # 1, N, C
+            query_pe = image_pe.reshape(1, -1, image_pe.shape[-1])
+            prompt = prompt_embeddings  # + prompt_pe
+            query = image_embeddings[b : (b + 1)]  # + query_pe
+            for block in self.prompt_blocks:
+                query, prompt = block(query, query_pe, prompt, prompt_pe)
+            image_embeddings_list.append(query[..., : image_embeddings.shape[-1]])
+            prompt_embeddings_list.append(prompt)
+        image_embeddings = torch.cat(image_embeddings_list, dim=0)
+        return image_embeddings
+
+
+class CrossAttenPromptBlock(nn.Module):
+    """
+    Self-attention block for prompt-based processing that handles both query and context features.
+    Supports different positional encoding strategies.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        init_values: float = 0.0,
+        first_block: bool = False,
+        pe: str = "normal",
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        self.first_block = first_block
+        self.pe = pe
+
+        # Attention components
+        self.norm1_x = nn.LayerNorm(dim)
+        self.norm1_x_after = nn.LayerNorm(dim)
+        self.attn_x = MemEffCrossAttention(dim, context_dim=dim, num_heads=num_heads, pe=pe)
+        self.ls1_x = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.norm1_context = nn.LayerNorm(dim)
+        self.attn_context = MemEffCrossAttention(dim, context_dim=dim, num_heads=num_heads, pe=pe)
+        self.ls1_context = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+
+        # MLP components
+        self.norm2_x = nn.LayerNorm(dim)
+        self.mlp_x = Mlp(dim, hidden_features=dim * 4)
+        self.ls2_x = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.norm2_context = nn.LayerNorm(dim)
+        self.mlp_context = Mlp(dim, hidden_features=dim * 4)
+        self.ls2_context = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+
+    def forward(self, x: Tensor, x_pe: Tensor, context: Tensor, context_pe: Tensor) -> Tuple[Tensor, Tensor]:
+        # Apply positional encoding if this is the first block and using normal PE
+        if self.pe == "normal" and self.first_block:
+            x = x + x_pe
+            context = context + context_pe
+
+        # Handle positional encoding concatenation if needed
+        if self.pe != "normal":
+            x = x + self.ls1_x(
+                self.attn_x(self.norm1_x(x), context=self.norm1_context(context), x_pe=x_pe, context_pe=context_pe)
+            )
+            context = context + self.ls1_context(
+                self.attn_context(
+                    self.norm1_context(context), context=self.norm1_x_after(x), x_pe=context_pe, context_pe=x_pe
+                )
+            )
+        else:
+            # Apply standard attention
+            x = x + self.ls1_x(self.attn_x(self.norm1_x(x), context=self.norm1_context(context)))
+            context = context + self.ls1_context(
+                self.attn_context(self.norm1_context(context), context=self.norm1_x_after(x))
+            )
+
+        # Apply MLP
+        x = x + self.ls2_x(self.mlp_x(self.norm2_x(x)))
+        context = context + self.ls2_context(self.mlp_context(self.norm2_context(context)))
+
+        return x, context

InfiniDepth/model/block/prompt_models/rope.py

@@ -0,0 +1,215 @@
+from typing import Dict, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+if torch.cuda.is_available():
+    acc_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
+else:
+    acc_dtype = torch.float16
+
+
+class PositionGetter:
+    """Generates and caches 2D spatial positions for patches in a grid.
+
+    This class efficiently manages the generation of spatial coordinates for patches
+    in a 2D grid, caching results to avoid redundant computations.
+
+    Attributes:
+        position_cache: Dictionary storing precomputed position tensors for different
+            grid dimensions.
+    """
+
+    def __init__(self):
+        """Initializes the position generator with an empty cache."""
+        self.position_cache: Dict[Tuple[int, int], torch.Tensor] = {}
+
+    def __call__(self, batch_size: int, height: int, width: int, device: torch.device) -> torch.Tensor:
+        """Generates spatial positions for a batch of patches.
+
+        Args:
+            batch_size: Number of samples in the batch.
+            height: Height of the grid in patches.
+            width: Width of the grid in patches.
+            device: Target device for the position tensor.
+
+        Returns:
+            Tensor of shape (batch_size, height*width, 2) containing y,x coordinates
+            for each position in the grid, repeated for each batch item.
+        """
+        if (height, width) not in self.position_cache:
+            y_coords = torch.arange(height, device=device)
+            x_coords = torch.arange(width, device=device)
+            positions = torch.cartesian_prod(y_coords, x_coords)
+            self.position_cache[height, width] = positions
+
+        cached_positions = self.position_cache[height, width]
+        return cached_positions.view(1, height * width, 2).expand(batch_size, -1, -1).clone()
+
+
+class RotaryPositionEmbedding2D(nn.Module):
+    """2D Rotary Position Embedding implementation.
+
+    This module applies rotary position embeddings to input tokens based on their
+    2D spatial positions. It handles the position-dependent rotation of features
+    separately for vertical and horizontal dimensions.
+
+    Args:
+        frequency: Base frequency for the position embeddings. Default: 100.0
+        scaling_factor: Scaling factor for frequency computation. Default: 1.0
+
+    Attributes:
+        base_frequency: Base frequency for computing position embeddings.
+        scaling_factor: Factor to scale the computed frequencies.
+        frequency_cache: Cache for storing precomputed frequency components.
+    """
+
+    def __init__(self, frequency: float = 100.0, scaling_factor: float = 1.0, feat_dim: int = 1024):
+        """Initializes the 2D RoPE module."""
+        super().__init__()
+        self.base_frequency = frequency
+        self.scaling_factor = scaling_factor
+        self.frequency_cache: Dict[Tuple, Tuple[torch.Tensor, torch.Tensor]] = {}
+        self.patch_size = 14
+        self.feat_dim = feat_dim
+
+    def _compute_frequency_components(
+        self, dim: int, seq_len: int, device: torch.device, dtype: torch.dtype
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Computes frequency components for rotary embeddings.
+
+        Args:
+            dim: Feature dimension (must be even).
+            seq_len: Maximum sequence length.
+            device: Target device for computations.
+            dtype: Data type for the computed tensors.
+
+        Returns:
+            Tuple of (cosine, sine) tensors for frequency components.
+        """
+        cache_key = (dim, seq_len, device, dtype)
+        if cache_key not in self.frequency_cache:
+            # Compute frequency bands
+            exponents = torch.arange(0, dim, 2, device=device).float() / dim
+            inv_freq = 1.0 / (self.base_frequency**exponents)
+
+            # Generate position-dependent frequencies
+            positions = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
+            angles = torch.einsum("i,j->ij", positions, inv_freq)
+
+            # Compute and cache frequency components
+            angles = angles.to(dtype)
+            angles = torch.cat((angles, angles), dim=-1)
+            cos_components = angles.cos().to(dtype)
+            sin_components = angles.sin().to(dtype)
+            self.frequency_cache[cache_key] = (cos_components, sin_components)
+
+        return self.frequency_cache[cache_key]
+
+    @staticmethod
+    def _rotate_features(x: torch.Tensor) -> torch.Tensor:
+        """Performs feature rotation by splitting and recombining feature dimensions.
+
+        Args:
+            x: Input tensor to rotate.
+
+        Returns:
+            Rotated feature tensor.
+        """
+        feature_dim = x.shape[-1]
+        x1, x2 = x[..., : feature_dim // 2], x[..., feature_dim // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+
+    def _apply_1d_rope(
+        self,
+        tokens: torch.Tensor,
+        positions: torch.Tensor,
+        cos_comp: torch.Tensor,
+        sin_comp: torch.Tensor,
+    ) -> torch.Tensor:
+        """Applies 1D rotary position embeddings along one dimension.
+
+        Args:
+            tokens: Input token features.
+            positions: Position indices.
+            cos_comp: Cosine components for rotation.
+            sin_comp: Sine components for rotation.
+
+        Returns:
+            Tokens with applied rotary position embeddings.
+        """
+        # Embed positions with frequency components
+        cos = F.embedding(positions, cos_comp)[:, None, :, :]
+        sin = F.embedding(positions, sin_comp)[:, None, :, :]
+
+        # Apply rotation
+        return (tokens * cos) + (self._rotate_features(tokens) * sin)
+
+    def _forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        """Applies 2D rotary position embeddings to input tokens.
+
+        Args:
+            tokens: Input tensor of shape (batch_size, n_heads, n_tokens, dim).
+                   The feature dimension (dim) must be divisible by 4.
+            positions: Position tensor of shape (batch_size, n_tokens, 2) containing
+                      the y and x coordinates for each token.
+
+        Returns:
+            Tensor of same shape as input with applied 2D rotary position embeddings.
+
+        Raises:
+            AssertionError: If input dimensions are invalid or positions are malformed.
+        """
+        # Validate inputs
+        assert tokens.size(-1) % 2 == 0, "Feature dimension must be even"
+        assert positions.ndim == 3 and positions.shape[-1] == 2, "Positions must have shape (batch_size, n_tokens, 2)"
+
+        # Compute feature dimension for each spatial direction
+        feature_dim = tokens.size(-1) // 2
+
+        # Get frequency components
+        max_position = int(positions.max()) + 1
+        cos_comp, sin_comp = self._compute_frequency_components(feature_dim, max_position, tokens.device, tokens.dtype)
+
+        # Split features for vertical and horizontal processing
+        vertical_features, horizontal_features = tokens.chunk(2, dim=-1)
+
+        # Apply RoPE separately for each dimension
+        vertical_features = self._apply_1d_rope(vertical_features, positions[..., 0], cos_comp, sin_comp)
+        horizontal_features = self._apply_1d_rope(horizontal_features, positions[..., 1], cos_comp, sin_comp)
+
+        # Combine processed features
+        return torch.cat((vertical_features, horizontal_features), dim=-1)
+
+    def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
+        """Positionally encode points that are normalized to [0,1]."""
+        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
+        max_position = int(coords.max()) + 1
+        cos_comp, sin_comp = self._compute_frequency_components(self.feat_dim, max_position, coords.device, acc_dtype)
+        vertical_cos = F.embedding(coords[..., 0], cos_comp)
+        vertical_sin = F.embedding(coords[..., 0], sin_comp)
+        horizontal_cos = F.embedding(coords[..., 1], cos_comp)
+        horizontal_sin = F.embedding(coords[..., 1], sin_comp)
+        # outputs d_1 x ... x d_n x C shape
+        return torch.cat((vertical_cos, vertical_sin, horizontal_cos, horizontal_sin), dim=-1)
+
+    def forward_encoding(self, size: Tuple[int, int], device: torch.device) -> torch.Tensor:
+        """Generate positional encoding for a grid of the specified size."""
+        height, width = size
+        y_coords = torch.arange(height, device=device) * (self.patch_size * 2) + self.patch_size - 1
+        x_coords = torch.arange(width, device=device) * (self.patch_size * 2) + self.patch_size - 1
+        positions = torch.cartesian_prod(y_coords, x_coords)  # h, w
+        max_position = int(positions.max()) + 1
+        cos_comp, sin_comp = self._compute_frequency_components(self.feat_dim, max_position, device, acc_dtype)
+        vertical_cos = F.embedding(positions[..., 0], cos_comp)
+        vertical_sin = F.embedding(positions[..., 0], sin_comp)
+        horizontal_cos = F.embedding(positions[..., 1], cos_comp)
+        horizontal_sin = F.embedding(positions[..., 1], sin_comp)
+        return (
+            torch.cat((vertical_cos, vertical_sin, horizontal_cos, horizontal_sin), dim=-1)
+            .reshape(height, width, -1)
+            .permute(2, 0, 1)
+        )
+
+    def forward(self, size: Tuple[int, int], device: torch.device) -> torch.Tensor:
+        return self.forward_encoding(size, device)

InfiniDepth/model/block/prompt_models/sam.py

@@ -0,0 +1,126 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Tuple, Type
+import torch
+from .utils.pe_utils import PositionEmbeddingRandom
+from .utils.transformer import TwoWayTransformer
+from torch import nn
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+class SAMPromptModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        transformer_dim: int,
+        mlp_dim: int = 2048,
+        num_heads: int = 8,
+        activation: Type[nn.Module] = nn.GELU,
+    ) -> None:
+        """
+        Predicts masks given an image and prompt embeddings, using a
+        transformer architecture.
+
+        Arguments:
+          transformer_dim (int): the channel dimension of the transformer
+          transformer (nn.Module): the transformer used to predict masks
+          num_multimask_outputs (int): the number of masks to predict
+            when disambiguating masks
+          activation (nn.Module): the type of activation to use when
+            upscaling masks
+          iou_head_depth (int): the depth of the MLP used to predict
+            mask quality
+          iou_head_hidden_dim (int): the hidden dimension of the MLP
+            used to predict mask quality
+        """
+        super().__init__()
+        self.transformer_dim = transformer_dim
+        self.transformer = TwoWayTransformer(
+            depth=2, embedding_dim=transformer_dim, num_heads=num_heads, mlp_dim=mlp_dim
+        )
+        self.pe_layer = PositionEmbeddingRandom(transformer_dim // 2)
+        self.depth2feature = nn.Sequential(
+            nn.Linear(1, transformer_dim // 2), nn.ReLU(True), nn.Linear(transformer_dim // 2, transformer_dim)
+        )
+
+    def forward(
+        self,
+        image_embeddings: torch.Tensor,
+        prompt_depth: torch.Tensor,
+        prompt_mask: torch.Tensor,
+        patch_h: int,
+        patch_w: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Predict masks given image and prompt embeddings.
+        Arguments:
+          image_embeddings (torch.Tensor): the embeddings from the image encoder
+          image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
+          sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
+          dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
+          multimask_output (bool): Whether to return multiple masks or a single
+            mask.
+        Returns:
+          torch.Tensor: batched predicted masks
+          torch.Tensor: batched predictions of mask quality
+        """
+        B, _, H, W = prompt_depth.shape
+        image_pe = self.pe_layer((patch_h, patch_w)).permute(1, 2, 0)  # CxHxW -> HxWxC
+
+        # prompt_embeddings_list = []
+        image_embeddings_list = []
+        for b in range(B):
+            valid_pts_num = (prompt_mask[b, 0] > 0.0).sum()
+            if valid_pts_num == 0:
+                image_embeddings_item = image_embeddings[b : (b + 1)]
+                image_embeddings_list.append(image_embeddings_item)
+                continue
+            sparse_depth_pos = (prompt_mask[b, 0] > 0.0).nonzero().float()
+            sparse_depth_pos[:, 0] = (sparse_depth_pos[:, 0] + 0.5) / H
+            sparse_depth_pos[:, 1] = (sparse_depth_pos[:, 1] + 0.5) / W
+            sparse_depth = prompt_depth[b, 0][prompt_mask[b, 0] > 0.0]
+            prompt_embeddings = self.depth2feature(sparse_depth[:, None])[None, ...]  # 1, N, C
+            prompt_pe = self.pe_layer._pe_encoding(sparse_depth_pos[None, :, [1, 0]])  # 1, N, C
+            prompt_embeddings_item, image_embeddings_item = self.transformer(
+                image_embeddings[b : (b + 1)],
+                image_pe.reshape(1, -1, image_pe.shape[-1]),
+                prompt_embeddings,
+                prompt_pe,
+            )
+            image_embeddings_list.append(image_embeddings_item)
+        image_embeddings = torch.cat(image_embeddings_list, dim=0)
+        return image_embeddings
+
+
+# # Lightly adapted from
+# # https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa
+# class MLP(nn.Module):
+#     def __init__(
+#         self,
+#         input_dim: int,
+#         hidden_dim: int,
+#         output_dim: int,
+#         num_layers: int,
+#         sigmoid_output: bool = False,
+#     ) -> None:
+#         super().__init__()
+#         self.num_layers = num_layers
+#         h = [hidden_dim] * (num_layers - 1)
+#         self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+#         self.sigmoid_output = sigmoid_output
+
+#     def forward(self, x):
+#         for i, layer in enumerate(self.layers):
+#             x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+#         if self.sigmoid_output:
+#             x = F.sigmoid(x)
+#         return x

InfiniDepth/model/block/prompt_models/selfattn.py

@@ -0,0 +1,289 @@
+from typing import Tuple
+import torch
+import torch.nn as nn
+from ..perceive_io import LayerScale, MemEffAttention, Mlp
+from .rope import RotaryPositionEmbedding2D
+from .utils.pe_utils import PositionEmbeddingRandom
+from torch import Tensor
+
+if torch.cuda.is_available():
+    acc_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
+else:
+    acc_dtype = torch.float16
+
+
+class SelfAttnPromptModel(nn.Module):
+    def __init__(
+        self,
+        transformer_dim: int = 1024,
+        num_blocks: int = 1,
+        num_heads: int = 4,
+        pe: str = "normal",
+        image_pe_method: str = "patch",  # image
+        **kwargs,
+    ) -> None:
+        """
+        Predicts masks given an image and prompt embeddings, using a
+        transformer architecture.
+
+        Arguments:
+          transformer_dim (int): the channel dimension of the transformer
+          transformer (nn.Module): the transformer used to predict masks
+          num_multimask_outputs (int): the number of masks to predict
+            when disambiguating masks
+          activation (nn.Module): the type of activation to use when
+            upscaling masks 
+          iou_head_depth (int): the depth of the MLP used to predict
+            mask quality
+          iou_head_hidden_dim (int): the hidden dimension of the MLP
+            used to predict mask quality
+        """
+        super().__init__()
+        self.pe = pe
+        pe_dim = transformer_dim // 2
+        if self.pe == "apply":
+            pe_dim = pe_dim // num_heads
+        self.pe_layer = PositionEmbeddingRandom(pe_dim, image_pe_method=image_pe_method)
+        self.prompt_blocks = nn.ModuleList(
+            [
+                SelfAttenPromptBlock(dim=transformer_dim, num_heads=num_heads, first_block=(i == 0), pe=pe)
+                for i in range(num_blocks)
+            ]
+        )
+        self.depth2feature = nn.Sequential(
+            nn.Linear(1, transformer_dim // 2),
+            nn.GELU(),
+            nn.Linear(transformer_dim // 2, transformer_dim),
+        )
+
+    def forward(
+        self,
+        image_embeddings: torch.Tensor,
+        prompt_depth: torch.Tensor,
+        prompt_mask: torch.Tensor,
+        patch_h: int,
+        patch_w: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Predict masks given image and prompt embeddings.
+        Arguments:
+          image_embeddings (torch.Tensor): the embeddings from the image encoder
+          image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
+          sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
+          dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
+          multimask_output (bool): Whether to return multiple masks or a single
+            mask.
+        Returns:
+          torch.Tensor: batched predicted masks
+          torch.Tensor: batched predictions of mask quality
+        """
+        B, _, H, W = prompt_depth.shape
+        image_pe = self.pe_layer((patch_h, patch_w)).permute(1, 2, 0)  # CxHxW -> HxWxC
+        prompt_embeddings_list = []
+        image_embeddings_list = []
+        for b in range(B):
+            valid_pts_num = (prompt_mask[b, 0] > 0.0).sum()
+            if valid_pts_num == 0:
+                image_embeddings_item = image_embeddings[b : (b + 1)]
+                image_embeddings_list.append(image_embeddings_item)
+                continue
+            sparse_depth_pos = (prompt_mask[b, 0] > 0.0).nonzero().float()
+            sparse_depth_pos[:, 0] = (sparse_depth_pos[:, 0] + 0.5) / H
+            sparse_depth_pos[:, 1] = (sparse_depth_pos[:, 1] + 0.5) / W
+            sparse_depth = prompt_depth[b, 0][prompt_mask[b, 0] > 0.0]
+            prompt_embeddings = self.depth2feature(sparse_depth[:, None])[None, ...]   # 1, N, C
+            prompt_pe = self.pe_layer._pe_encoding(sparse_depth_pos[None, :, [1, 0]])  # 1, N, C
+            query_pe = image_pe.reshape(1, -1, image_pe.shape[-1])
+            prompt = prompt_embeddings  # + prompt_pe
+            query = image_embeddings[b : (b + 1)]  # + query_pe
+            with torch.autocast("cuda", enabled=True, dtype=acc_dtype):
+                for block in self.prompt_blocks:
+                    query, prompt = block(query, query_pe, prompt, prompt_pe)
+            image_embeddings_list.append(query[..., : image_embeddings.shape[-1]])
+            prompt_embeddings_list.append(prompt)
+        image_embeddings = torch.cat(image_embeddings_list, dim=0)
+        return image_embeddings
+
+
+class SelfAttnRopePromptModel(nn.Module):
+    def __init__(
+        self,
+        transformer_dim: int = 1024,
+        num_blocks: int = 1,
+        num_heads: int = 4,
+        pe: str = "normal",
+        image_pe_method: str = "patch",  # image
+        **kwargs,
+    ) -> None:
+        """
+        Predicts masks given an image and prompt embeddings, using a
+        transformer architecture.
+
+        Arguments:
+          transformer_dim (int): the channel dimension of the transformer
+          transformer (nn.Module): the transformer used to predict masks
+          num_multimask_outputs (int): the number of masks to predict
+            when disambiguating masks
+          activation (nn.Module): the type of activation to use when
+            upscaling masks
+          iou_head_depth (int): the depth of the MLP used to predict
+            mask quality
+          iou_head_hidden_dim (int): the hidden dimension of the MLP
+            used to predict mask quality
+        """
+        super().__init__()
+        self.pe = pe
+        pe_dim = transformer_dim // 2
+        if self.pe == "apply":
+            pe_dim = pe_dim // num_heads
+        if self.pe.startswith("rope"):
+            self.pe_layer = RotaryPositionEmbedding2D(
+                frequency=float(self.pe.split("rope")[1]), feat_dim=pe_dim // num_heads
+            )
+        else:
+            self.pe_layer = PositionEmbeddingRandom(pe_dim, image_pe_method=image_pe_method)
+        self.prompt_blocks = nn.ModuleList(
+            [
+                SelfAttenPromptBlock(
+                    dim=transformer_dim, num_heads=num_heads, first_block=(i == 0), pe=pe, use_sep=False
+                )
+                for i in range(num_blocks)
+            ]
+        )
+        self.depth2feature = nn.Sequential(
+            nn.Linear(1, transformer_dim // 2),
+            nn.GELU(),
+            nn.Linear(transformer_dim // 2, transformer_dim),
+        )
+
+    def forward(
+        self,
+        image_embeddings: torch.Tensor,
+        prompt_depth: torch.Tensor,
+        prompt_mask: torch.Tensor,
+        patch_h: int,
+        patch_w: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Predict masks given image and prompt embeddings.
+        Arguments:
+          image_embeddings (torch.Tensor): the embeddings from the image encoder
+          image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
+          sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
+          dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
+          multimask_output (bool): Whether to return multiple masks or a single
+            mask.
+        Returns:
+          torch.Tensor: batched predicted masks
+          torch.Tensor: batched predictions of mask quality
+        """
+        B, _, H, W = prompt_depth.shape
+        image_pe = self.pe_layer((patch_h, patch_w), device=prompt_depth.device).permute(1, 2, 0)  # CxHxW -> HxWxC
+        prompt_embeddings_list = []
+        image_embeddings_list = []
+        for b in range(B):
+            valid_pts_num = (prompt_mask[b, 0] > 0.0).sum()
+            if valid_pts_num == 0:
+                image_embeddings_item = image_embeddings[b : (b + 1)]
+                image_embeddings_list.append(image_embeddings_item)
+                continue
+            sparse_depth_pos = (prompt_mask[b, 0] > 0.0).nonzero().int()
+            sparse_depth_pos[:, 0] = sparse_depth_pos[:, 0] * 2
+            sparse_depth_pos[:, 1] = sparse_depth_pos[:, 1] * 2
+            sparse_depth = prompt_depth[b, 0][prompt_mask[b, 0] > 0.0]
+            prompt_embeddings = self.depth2feature(sparse_depth[:, None])[None, ...]  # 1, N, C
+            prompt_pe = self.pe_layer._pe_encoding(sparse_depth_pos[None])  # 1, N, C
+            query_pe = image_pe.reshape(1, -1, image_pe.shape[-1])
+            prompt = prompt_embeddings  # + prompt_pe
+            query = image_embeddings[b : (b + 1)]  # + query_pe
+            with torch.autocast("cuda", enabled=True, dtype=acc_dtype):
+                for block in self.prompt_blocks:
+                    query, prompt = block(query, query_pe, prompt, prompt_pe)
+            image_embeddings_list.append(query[..., : image_embeddings.shape[-1]])
+            prompt_embeddings_list.append(prompt)
+        image_embeddings = torch.cat(image_embeddings_list, dim=0)
+        return image_embeddings
+    
+
+class SelfAttenPromptBlock(nn.Module):
+    """
+    Self-attention block for prompt-based processing that handles both query and context features.
+    Supports different positional encoding strategies.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        init_values: float = 0.0,
+        first_block: bool = False,
+        pe: str = "normal",
+        use_sep: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        self.first_block = first_block
+        self.pe = pe
+
+        # Attention components
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = MemEffAttention(dim, num_heads=num_heads, pe=pe)
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+
+        # MLP components
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = Mlp(dim, hidden_features=dim * 4)
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+
+        # Separator token for concatenating query and context
+        pe_dim = dim
+        self.use_sep = use_sep
+        if use_sep:
+            self.sep = nn.Parameter(torch.randn(1, 1, pe_dim))
+        else:
+            self.sep = None
+
+        # Special separator for positional encoding if needed
+        if self.use_sep:
+            if self.pe != "normal":
+                self.sep_pe = nn.Parameter(torch.randn(1, 1, pe_dim))
+            else:
+                self.sep_pe = None
+
+    def forward(self, x: Tensor, x_pe: Tensor, context: Tensor, context_pe: Tensor) -> Tuple[Tensor, Tensor]:
+        # Apply positional encoding if this is the first block and using normal PE
+        if self.pe == "normal" and self.first_block:
+            x = x + x_pe
+            context = context + context_pe
+
+        # Record original sequence lengths
+        x_len, context_len = x.shape[1], context.shape[1]
+
+        # Concatenate query, separator token, and context
+        if self.use_sep:
+            x = torch.cat([x, self.sep, context], dim=1)
+        else:
+            x = torch.cat([x, context], dim=1)
+
+        # Handle positional encoding concatenation if needed
+        if self.pe != "normal":
+            if self.use_sep:
+                x_pe = torch.cat([x_pe, self.sep_pe, context_pe], dim=1)
+            else:
+                x_pe = torch.cat([x_pe, context_pe], dim=1)
+            x = x + self.ls1(self.attn(self.norm1(x), x_pe))
+        else:
+            # Apply standard attention
+            x = x + self.ls1(self.attn(self.norm1(x)))
+
+        # Apply MLP
+        x = x + self.ls2(self.mlp(self.norm2(x)))
+
+        # Split back into query and context
+        query = x[:, :x_len, :]
+        if self.use_sep:
+            context = x[:, x_len + 1 : x_len + 1 + context_len, :]
+        else:
+            context = x[:, x_len : x_len + context_len, :]
+
+        return query, context

InfiniDepth/model/block/prompt_models/utils/__init__.py

@@ -0,0 +1 @@
+"""Prompt model utility modules."""

InfiniDepth/model/block/prompt_models/utils/pe_utils.py

@@ -0,0 +1,72 @@
+from typing import Any, Optional, Tuple
+import numpy as np
+import torch
+from torch import nn
+
+
+class PositionEmbeddingRandom(nn.Module):
+    """
+    Positional encoding using random spatial frequencies.
+    """
+
+    patch_size = 14
+
+    def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None, image_pe_method: str = "patch") -> None:
+        super().__init__()
+        if scale is None or scale <= 0.0:
+            scale = 1.0
+        self.register_buffer(
+            "positional_encoding_gaussian_matrix",
+            scale * torch.randn((2, num_pos_feats)),
+        )
+
+        self.image_pe_method = image_pe_method
+        if self.image_pe_method == "image":
+            # self.patch_embed = nn.Conv2d(num_pos_feats*2, num_pos_feats*2, kernel_size=self.patch_size, stride=self.patch_size)
+            self.patch_embed = nn.Sequential(
+                nn.Conv2d(num_pos_feats * 2, num_pos_feats // 2, kernel_size=2, stride=2),
+                nn.ReLU(),
+                nn.Conv2d(
+                    num_pos_feats // 2,
+                    num_pos_feats * 2,
+                    kernel_size=self.patch_size // 2,
+                    stride=self.patch_size // 2,
+                ),
+            )
+
+    def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
+        """Positionally encode points that are normalized to [0,1]."""
+        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
+        coords = 2 * coords - 1
+        coords = coords @ self.positional_encoding_gaussian_matrix
+        coords = 2 * np.pi * coords
+        # outputs d_1 x ... x d_n x C shape
+        return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
+
+    def forward_encoding(self, size: Tuple[int, int]) -> torch.Tensor:
+        """Generate positional encoding for a grid of the specified size."""
+        h, w = size
+        device: Any = self.positional_encoding_gaussian_matrix.device
+        grid = torch.ones((h, w), device=device, dtype=torch.float32)
+        y_embed = grid.cumsum(dim=0) - 0.5
+        x_embed = grid.cumsum(dim=1) - 0.5
+        y_embed = y_embed / h
+        x_embed = x_embed / w
+        pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))  # HxWx2 -> HxWxC
+        return pe.permute(2, 0, 1)  # C x H x W
+
+    def forward(self, size: Tuple[int, int]) -> torch.Tensor:
+
+        if self.image_pe_method == "patch":
+            return self.forward_encoding(size)
+        elif self.image_pe_method == "image":
+            pe_encoding = self.forward_encoding(size)
+            pe_encoding_high = self.forward_encoding((size[0] * self.patch_size, size[1] * self.patch_size))
+            return pe_encoding + self.patch_embed(pe_encoding_high[None])[0]
+
+    def forward_with_coords(self, coords_input: torch.Tensor, image_size: Tuple[int, int]) -> torch.Tensor:
+        """Positionally encode points that are not normalized to [0,1]."""
+        coords = coords_input.clone()
+        coords[:, :, 0] = coords[:, :, 0] / image_size[1]
+        coords[:, :, 1] = coords[:, :, 1] / image_size[0]
+        return self._pe_encoding(coords.to(torch.float))  # B x N x C

InfiniDepth/model/block/prompt_models/utils/transformer.py

@@ -0,0 +1,250 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Tuple, Type
+import torch
+from einops import rearrange
+from ...common import MLPBlock
+from .....utils.logger import Log
+from torch import Tensor, nn
+
+try:
+    from xformers.ops import memory_efficient_attention
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    Log.warning("xFormers not available")
+    XFORMERS_AVAILABLE = False
+
+
+class TwoWayTransformer(nn.Module):
+    def __init__(
+        self,
+        depth: int,
+        embedding_dim: int,
+        num_heads: int,
+        mlp_dim: int,
+        activation: Type[nn.Module] = nn.ReLU,
+        attention_downsample_rate: int = 2,
+    ) -> None:
+        """
+        A transformer decoder that attends to an input image using
+        queries whose positional embedding is supplied.
+
+        Args:
+          depth (int): number of layers in the transformer
+          embedding_dim (int): the channel dimension for the input embeddings
+          num_heads (int): the number of heads for multihead attention. Must
+            divide embedding_dim
+          mlp_dim (int): the channel dimension internal to the MLP block
+          activation (nn.Module): the activation to use in the MLP block
+        """
+        super().__init__()
+        self.depth = depth
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.mlp_dim = mlp_dim
+        self.layers = nn.ModuleList()
+
+        for i in range(depth):
+            self.layers.append(
+                TwoWayAttentionBlock(
+                    embedding_dim=embedding_dim,
+                    num_heads=num_heads,
+                    mlp_dim=mlp_dim,
+                    activation=activation,
+                    attention_downsample_rate=attention_downsample_rate,
+                    skip_first_layer_pe=(i == 0),
+                )
+            )
+
+        self.final_attn_token_to_image = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)
+        self.norm_final_attn = nn.LayerNorm(embedding_dim)
+
+    def forward(
+        self, image_embedding: Tensor, image_pe: Tensor, point_embedding: Tensor, point_pe: Tensor
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Args:
+          image_embedding (torch.Tensor): image to attend to. Should be shape
+            B x embedding_dim x h x w for any h and w.
+          image_pe (torch.Tensor): the positional encoding to add to the image. Must
+            have the same shape as image_embedding.
+          point_embedding (torch.Tensor): the embedding to add to the query points.
+            Must have shape B x N_points x embedding_dim for any N_points.
+
+        Returns:
+          torch.Tensor: the processed point_embedding
+          torch.Tensor: the processed image_embedding
+        """
+        # Prepare queries
+        queries = point_embedding
+        keys = image_embedding
+
+        # Apply transformer blocks and final layernorm
+        for layer in self.layers:
+            queries, keys = layer(
+                queries=queries,
+                keys=keys,
+                query_pe=point_pe,
+                key_pe=image_pe,
+            )
+        # queries become keys-aware; keys become queries-aware
+
+        # Apply the final attention layer from the points to the image
+        q = queries + point_embedding
+        k = keys + image_pe
+        attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
+        queries = queries + attn_out
+        queries = self.norm_final_attn(queries)
+
+        return queries, keys
+
+
+class TwoWayAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_heads: int,
+        mlp_dim: int = 2048,
+        activation: Type[nn.Module] = nn.ReLU,
+        attention_downsample_rate: int = 2,
+        skip_first_layer_pe: bool = False,
+    ) -> None:
+        """
+        A transformer block with four layers: (1) self-attention of sparse
+        inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp
+        block on sparse inputs, and (4) cross attention of dense inputs to sparse
+        inputs.
+
+        Arguments:
+          embedding_dim (int): the channel dimension of the embeddings
+          num_heads (int): the number of heads in the attention layers
+          mlp_dim (int): the hidden dimension of the mlp block
+          activation (nn.Module): the activation of the mlp block
+          skip_first_layer_pe (bool): skip the PE on the first layer
+        """
+        super().__init__()
+        self.self_attn = MemEffAttention(embedding_dim, num_heads)
+        self.norm1 = nn.LayerNorm(embedding_dim)
+
+        self.cross_attn_token_to_image = MemEffAttention(
+            embedding_dim, num_heads, downsample_rate=attention_downsample_rate
+        )
+        self.norm2 = nn.LayerNorm(embedding_dim)
+
+        self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
+        self.norm3 = nn.LayerNorm(embedding_dim)
+
+        self.norm4 = nn.LayerNorm(embedding_dim)
+        self.cross_attn_image_to_token = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)
+
+        self.skip_first_layer_pe = skip_first_layer_pe
+
+    def forward(self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor) -> Tuple[Tensor, Tensor]:
+        # Self attention block
+        #
+        if self.skip_first_layer_pe:
+            queries = self.self_attn(q=queries, k=queries, v=queries)
+        else:
+            q = queries + query_pe
+            attn_out = self.self_attn(q=q, k=q, v=queries)
+            queries = queries + attn_out
+        queries = self.norm1(queries)
+
+        # Cross attention block, tokens attending to image embedding
+        q = queries + query_pe
+        k = keys + key_pe
+        attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
+        queries = queries + attn_out
+        queries = self.norm2(queries)
+
+        # MLP block
+        mlp_out = self.mlp(queries)
+        queries = queries + mlp_out
+        queries = self.norm3(queries)
+
+        # Cross attention block, image embedding attending to tokens
+        q = queries + query_pe
+        k = keys + key_pe
+        attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
+        keys = keys + attn_out
+        keys = self.norm4(keys)
+
+        return queries, keys
+
+
+class Attention(nn.Module):
+    """
+    An attention layer that allows for downscaling the size of the embedding
+    after projection to queries, keys, and values.
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_heads: int,
+        downsample_rate: int = 1,
+    ) -> None:
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.internal_dim = embedding_dim // downsample_rate
+        self.num_heads = num_heads
+        assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim."
+
+        self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
+        self.k_proj = nn.Linear(embedding_dim, self.internal_dim)
+        self.v_proj = nn.Linear(embedding_dim, self.internal_dim)
+        self.out_proj = nn.Linear(self.internal_dim, embedding_dim)
+
+    def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
+        b, n, c = x.shape
+        x = x.reshape(b, n, num_heads, c // num_heads)
+        return x.transpose(1, 2)  # B x N_heads x N_tokens x C_per_head
+
+    def _recombine_heads(self, x: Tensor) -> Tensor:
+        b, n_heads, n_tokens, c_per_head = x.shape
+        x = x.transpose(1, 2)
+        return x.reshape(b, n_tokens, n_heads * c_per_head)  # B x N_tokens x C
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        # Input projections
+        q = self.q_proj(q)
+        k = self.k_proj(k)
+        v = self.v_proj(v)
+
+        # Separate into heads
+        q = self._separate_heads(q, self.num_heads)
+        k = self._separate_heads(k, self.num_heads)
+        v = self._separate_heads(v, self.num_heads)
+
+        # Attention
+        _, _, _, c_per_head = q.shape
+        attn = q @ k.permute(0, 1, 3, 2)  # B x N_heads x N_tokens x N_tokens
+        attn = attn / math.sqrt(c_per_head)
+        attn = torch.softmax(attn, dim=-1)
+
+        # Get output
+        out = attn @ v
+        out = self._recombine_heads(out)
+        out = self.out_proj(out)
+
+        return out
+
+
+class MemEffAttention(Attention):
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        q = self.q_proj(q)
+        k = self.k_proj(k)
+        v = self.v_proj(v)
+        q = rearrange(q, "b n (m c) -> b n m c", m=self.num_heads)
+        k = rearrange(k, "b n (m c) -> b n m c", m=self.num_heads)
+        v = rearrange(v, "b n (m c) -> b n m c", m=self.num_heads)
+        x = memory_efficient_attention(q, k, v, attn_bias=None)
+        x = rearrange(x, "b n m c -> b n (m c)")
+        x = self.out_proj(x)
+        return x

InfiniDepth/model/block/rope.py

@@ -0,0 +1,69 @@
+from typing import Tuple
+import torch
+
+
+def precompute_freqs_cis_2d(dim: int, height: int, width:int, theta: float = 10000.0, scale=16.0):
+    # assert  H * H == end
+    # flat_patch_pos = torch.linspace(-1, 1, end) # N = end
+    x_pos = torch.linspace(0, scale, width)
+    y_pos = torch.linspace(0, scale, height)
+    y_pos, x_pos = torch.meshgrid(y_pos, x_pos, indexing="ij")
+    y_pos = y_pos.reshape(-1)
+    x_pos = x_pos.reshape(-1)
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) # Hc/4
+    x_freqs = torch.outer(x_pos, freqs).float() # N Hc/4
+    y_freqs = torch.outer(y_pos, freqs).float() # N Hc/4
+    x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)
+    y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)
+    freqs_cis = torch.cat([x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1) # N,Hc/4,2
+    freqs_cis = freqs_cis.reshape(height*width, -1)
+    return freqs_cis
+
+def precompute_freqs_cis_ex2d(dim: int, height: int, width:int, theta: float = 10000.0, scale=1.0):
+    if isinstance(scale, float):
+        scale = (scale, scale)
+    x_pos = torch.linspace(0, height*scale[0], width)
+    y_pos = torch.linspace(0, width*scale[1], height)
+    y_pos, x_pos = torch.meshgrid(y_pos, x_pos, indexing="ij")
+    y_pos = y_pos.reshape(-1)
+    x_pos = x_pos.reshape(-1)
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) # Hc/4
+    x_freqs = torch.outer(x_pos, freqs).float() # N Hc/4
+    y_freqs = torch.outer(y_pos, freqs).float() # N Hc/4
+    x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)
+    y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)
+    freqs_cis = torch.cat([x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1) # N,Hc/4,2
+    freqs_cis = freqs_cis.reshape(height*width, -1)
+    return freqs_cis
+
+
+def apply_rotary_emb(
+        xq: torch.Tensor,
+        xk: torch.Tensor,
+        freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    freqs_cis = freqs_cis[None, None, :, :]
+    # xq : B N H Hc
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # B N H Hc/2
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) # B, N, H, Hc
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+def apply_rotary_emb_crossattention(
+        xq: torch.Tensor,
+        xk: torch.Tensor,
+        yk: torch.Tensor,
+        freqs_cis1: torch.Tensor,
+        freqs_cis2: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    freqs_cis1 = freqs_cis1[None, None, :, :]
+    freqs_cis2 = freqs_cis2[None, None, :, :]
+    # xq : B N H Hc
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # B N H Hc/2
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    yk_ = torch.view_as_complex(yk.float().reshape(*yk.shape[:-1], -1, 2))
+    xq_out = torch.view_as_real(xq_ * freqs_cis1).flatten(3) # B, N, H, Hc
+    xk_out = torch.view_as_real(xk_ * freqs_cis1).flatten(3)
+    yk_out = torch.view_as_real(yk_ * freqs_cis2).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk), yk_out.type_as(yk)

InfiniDepth/model/block/torchhub/README.md

@@ -0,0 +1,3 @@
+# Local PyTorch Hub
+
+This directory is for loading the DINOv2 encoder locally in case of no Internet connection.

InfiniDepth/model/block/torchhub/dinov3/.docstr.yaml

@@ -0,0 +1,6 @@
+paths:
+    - dinov3
+exclude: dinov3/tests
+skip_init: True
+skip_private: True
+fail_under: 0

InfiniDepth/model/block/torchhub/dinov3/.github/workflows/lint.yaml

@@ -0,0 +1,47 @@
+name: Lint
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  run-linters:
+    name: Run linters
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.11
+          cache: 'pip'
+          cache-dependency-path: '**/requirements*.txt'
+      - name: Install Python (development) dependencies
+        run: |
+          pip install -r requirements-dev.txt
+      - name: Run ruff (linter)
+        run: |
+          ruff check dinov3
+      - name: Run ruff (formatter)
+        if: always()
+        run: |
+          ruff format --diff dinov3
+      - name: Report docstring coverage
+        if: always()
+        run: |
+          docstr-coverage dinov3
+      - name: Run mypy
+        if: always()
+        run: |
+          mypy --txt-report .
+          [ -f index.txt ] && cat index.txt
+      - name: Run pylint
+        if: always()
+        run: |
+          pylint --exit-zero dinov3

InfiniDepth/model/block/torchhub/dinov3/.gitignore

@@ -0,0 +1,18 @@
+build/
+dist/
+*.egg-info/
+**/__pycache__/
+
+**/.ipynb_checkpoints
+**/.ipynb_checkpoints/**
+
+**/notebooks
+
+# Ignore shell scripts
+*.sh
+
+# Ignore swap files
+*.swp
+
+# Ignore vscode directory
+.vscode/

InfiniDepth/model/block/torchhub/dinov3/CODE_OF_CONDUCT.md

@@ -0,0 +1,80 @@
+# Code of Conduct
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to make participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, sex characteristics, gender identity and expression,
+level of experience, education, socio-economic status, nationality, personal
+appearance, race, religion, or sexual identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies within all project spaces, and it also applies when
+an individual is representing the project or its community in public spaces.
+Examples of representing a project or community include using an official
+project e-mail address, posting via an official social media account, or acting
+as an appointed representative at an online or offline event. Representation of
+a project may be further defined and clarified by project maintainers.
+
+This Code of Conduct also applies outside the project spaces when there is a
+reasonable belief that an individual's behavior may have a negative impact on
+the project or its community.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at <opensource-conduct@meta.com>. All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
+available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see
+https://www.contributor-covenant.org/faq

InfiniDepth/model/block/torchhub/dinov3/CONTRIBUTING.md

@@ -0,0 +1,31 @@
+# Contributing to DINOv3
+We want to make contributing to this project as easy and transparent as
+possible.
+
+## Pull Requests
+We actively welcome your pull requests.
+
+1. Fork the repo and create your branch from `main`.
+2. If you've added code that should be tested, add tests.
+3. If you've changed APIs, update the documentation.
+4. Ensure the test suite passes.
+5. Make sure your code lints.
+6. If you haven't already, complete the Contributor License Agreement ("CLA").
+
+## Contributor License Agreement ("CLA")
+In order to accept your pull request, we need you to submit a CLA. You only need
+to do this once to work on any of Meta's open source projects.
+
+Complete your CLA here: <https://code.facebook.com/cla>
+
+## Issues
+We use GitHub issues to track public bugs. Please ensure your description is
+clear and has sufficient instructions to be able to reproduce the issue.
+
+Meta has a [bounty program](https://www.facebook.com/whitehat/) for the safe
+disclosure of security bugs. In those cases, please go through the process
+outlined on that page and do not file a public issue.
+
+## License
+By contributing to DINOv3, you agree that your contributions will be licensed
+under the LICENSE.md file in the root directory of this source tree.

InfiniDepth/model/block/torchhub/dinov3/LICENSE.md

@@ -0,0 +1,66 @@
+# DINOv3 License
+
+*Last Updated: August 19, 2025*
+
+**“Agreement”** means the terms and conditions for use, reproduction, distribution and modification of the DINO Materials set forth herein.
+
+**“DINO Materials”** means, collectively, Documentation and the models, software and algorithms, including machine-learning model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code, and other elements of the foregoing distributed by Meta and made available under this Agreement.
+
+**“Documentation”** means the specifications, manuals and documentation accompanying
+DINO Materials distributed by Meta.
+
+**“Licensee”** or **“you”** means you, or your employer or any other person or entity (if you are entering into this Agreement on such person or entity’s behalf), of the age required under applicable laws, rules or regulations to provide legal consent and that has legal authority to bind your employer or such other person or entity if you are entering in this Agreement on their behalf.
+
+**“Meta”** or **“we”** means Meta Platforms Ireland Limited (if you are located in or, if you are an entity, your principal place of business is in the EEA or Switzerland) or Meta Platforms, Inc. (if you are located outside of the EEA or Switzerland).
+
+**“Sanctions”** means any economic or trade sanctions or restrictions administered or enforced by the United States (including the Office of Foreign Assets Control of the U.S. Department of the Treasury (“OFAC”), the U.S. Department of State and the U.S. Department of Commerce), the United Nations, the European Union, or the United Kingdom.
+
+**“Trade Controls”** means any of the following: Sanctions and applicable export and import controls.
+
+By clicking “I Accept” below or by using or distributing any portion or element of the DINO Materials, you agree to be bound by this Agreement.
+
+## 1. License Rights and Redistribution.
+
+a. <ins>Grant of Rights</ins>. You are granted a non-exclusive, worldwide, non-transferable and royalty-free limited license under Meta’s intellectual property or other rights owned by Meta embodied in the DINO Materials to use, reproduce, distribute, copy, create derivative works of, and make modifications to the DINO Materials.
+
+b. <ins>Redistribution and Use</ins>.
+
+i. Distribution of DINO Materials, and any derivative works thereof, are subject to the terms of this Agreement. If you distribute or make the DINO Materials, or any derivative works thereof, available to a third party, you may only do so under the terms of this Agreement and you shall provide a copy of this Agreement with any such DINO Materials.
+
+ii.  If you submit for publication the results of research you perform on, using, or otherwise in connection with DINO Materials, you must acknowledge the use of DINO Materials in your publication.
+
+iii. Your use of the DINO Materials must comply with applicable laws and regulations, including Trade Control Laws and applicable privacy and data protection laws.
+
+iv. Your use of the DINO Materials will not involve or encourage others to reverse engineer, decompile or discover the underlying components of the DINO Materials.
+
+v. You are not the target of Trade Controls and your use of DINO Materials must comply with Trade Controls. You agree not to use, or permit others to use, DINO Materials for any activities subject to the International Traffic in Arms Regulations (ITAR) or end uses prohibited by Trade Controls, including those related to military or warfare purposes, nuclear industries or applications, espionage, or the development or use of guns or illegal weapons.
+
+## 2. User Support.
+
+Your use of the DINO Materials is done at your own discretion; Meta does not process any information nor provide any service in relation to such use.  Meta is under no obligation to provide any support services for the DINO Materials. Any support provided is “as is”, “with all faults”, and without warranty of any kind.
+
+## 3. Disclaimer of Warranty.
+
+UNLESS REQUIRED BY APPLICABLE LAW, THE DINO MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” BASIS, WITHOUT WARRANTIES OF ANY KIND, AND META DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING THE DINO MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE DINO MATERIALS AND ANY OUTPUT AND RESULTS.
+
+## 4. Limitation of Liability.
+
+IN NO EVENT WILL META OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY DIRECT OR INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF META OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+
+## 5. Intellectual Property.
+
+a. Subject to Meta’s ownership of DINO Materials and derivatives made by or for Meta, with respect to any derivative works and modifications of the DINO Materials that are made by you, as between you and Meta, you are and will be the owner of such derivative works and modifications.
+
+b. If you institute litigation or other proceedings against Meta or any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the DINO Materials, outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by you, then any licenses granted to you under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Meta from and against any claim by any third party arising out of or related to your use or distribution of the DINO Materials.
+
+## 6. Term and Termination.
+
+The term of this Agreement will commence upon your acceptance of this Agreement or access to the DINO Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Meta may terminate this Agreement if you are in breach of any term or condition of this Agreement. Upon termination of this Agreement, you shall delete and cease use of the DINO Materials. Sections 3, 4 and 7 shall survive the termination of this Agreement.
+
+## 7. Governing Law and Jurisdiction.
+
+This Agreement will be governed and construed under the laws of the State of California without regard to choice of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The courts of California shall have exclusive jurisdiction of any dispute arising out of this Agreement.
+
+## 8. Modifications and Amendments.
+
+Meta may modify this Agreement from time to time; provided that they are similar in spirit to the current version of the Agreement, but may differ in detail to address new problems or concerns. All such changes will be effective immediately. Your continued use of the DINO Materials after any modification to this Agreement constitutes your agreement to such modification. Except as provided in this Agreement, no modification or addition to any provision of this Agreement will be binding unless it is in writing and signed by an authorized representative of both you and Meta.

InfiniDepth/model/block/torchhub/dinov3/MODEL_CARD.md

@@ -0,0 +1,432 @@
+# Model Card for DINOv3
+
+DINOv3 is a family of versatile vision foundation models that outperforms the specialized state of the art across a broad range of settings, without fine-tuning. DINOv3 produces high-quality dense features that achieve outstanding performance on various vision tasks, significantly surpassing previous self- and weakly-supervised foundation models.
+
+## Model Details
+
+These are Vision Transformer and ConvNeXt models trained following the method described in the DINOv3 paper. 12 models are provided:
+
+- 10 models pretrained on web data (LVD-1689M dataset)
+  - 1 ViT-7B trained from scratch,
+  - 5 ViT-S/S+/B/L/H+ models distilled from the ViT-7B,
+  - 4 ConvNeXt-{T/S/B/L} models distilled from the ViT-7B,
+- 2 models pretrained on satellite data (SAT-493M dataset)
+  - 1 ViT-7B trained from scratch
+  - 1 ViT-L distilled from the ViT-7B
+
+
+Each Transformer-based model takes an image as input and returns a class token, patch tokens (and register tokens). These models follow a ViT architecture, with a patch size of 16. For a 224x224 image, this results in 1 class token + 4 register tokens + 196 patch tokens = 201 tokens (for DINOv2 with registers this resulted in 1 + 4 + 256 = 261 tokens).
+
+The models can accept larger images provided the image shapes are multiples of the patch size (16). If this condition is not verified, the model will crop to the closest smaller multiple of the patch size.
+
+### Model Description
+
+- **Developed by:** Meta AI
+- **Model type:** Vision Transformer, ConvNeXt
+- **License:** [DINOv3 License](https://ai.meta.com/resources/models-and-libraries/dinov3-license/)
+
+### Model Sources
+
+- **Repository:** [https://github.com/facebookresearch/dinov3](https://github.com/facebookresearch/dinov3)
+- **Paper:** [https://arxiv.org/abs/2508.10104](https://arxiv.org/abs/2508.10104)
+
+## Uses
+
+The models are vision backbones providing multi-purpose features for downstream tasks.
+
+### Direct Use
+
+The models can be used without fine-tuning, with downstream classifiers as simple as linear layers, to obtain competitive results:
+
+- on image classification, using k-NN classifiers on the class token
+- on image classification, with logistic regression classifiers applied on the class token
+- on image classification, with a linear layer applied on the class token and the average of the patch tokens
+- on image retrieval using nearest neighbors
+- on geometric and semantic 3D keypoint correspondances
+- on depth estimation, semantic segmentation, using linear layers
+- on unsupervised object discovery
+- on video segmentation tracking
+- on video classification, using a small 4-layer attentive probe
+
+### Downstream Use
+
+While fine-tuning the models can yield some gains, it is recommended to keep this option as a last resort: the frozen features are expected to provide good performance out-of-the-box.
+
+## Bias, Risks, and Limitations
+
+Compared to DINOv2 and SEERv2, DINOv3 delivers somewhat consistent performance across income categories on geographical fairness and diversity, although with a notable performance drop in the low-income bucket compared to the highest-income bucket.
+
+DINOv3 also achieves relatively good scores across different regions, improving over its predecessor DINOv2. However, a relative difference is still observed between Europe and Africa.
+
+### Recommendations
+
+Fine-tuning is expected to increase the biases in the features produced by the model as they will be tuned to the fine-tuning labels.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+```python
+import torch
+
+model = torch.hub.load(
+    repo_or_dir='facebookresearch/dinov3',
+    model='<MODEL_NAME>',
+    weights='<PATH/OR/URL/TO/CHECKPOINT>',
+)
+
+# where MODEL_NAME can be one of:
+# - dinov3_vits16
+# - dinov3_vits16plus
+# - dinov3_vitb16
+# - dinov3_vitl16
+# - dinov3_vith16plus
+# - dinov3_vit7b16
+# - dinov3_convnext_tiny
+# - dinov3_convnext_small
+# - dinov3_convnext_base
+# - dinov3_convnext_large
+
+# For instance
+dinov3_vits16 = torch.hub.load(
+    repo_or_dir='facebookresearch/dinov3',
+    model='dinov3_vits16',
+    weights='<PATH/OR/URL/TO/DINOV3/VITS16/LVD1689M/CHECKPOINT>',
+)
+```
+
+## Training Details
+
+### Training Data
+
+- Web dataset (LVD-1689M): a curated dataset of 1,689 millions of images extracted from a large data
+pool of 17 billions web images collected from public posts on Instagram
+
+- Satellite dataset (SAT-493M): a dataset of 493 millions of 512x512 images sampled randomly from Maxar RGB ortho-rectified imagery at 0.6 meter resolution
+
+### Training Procedure
+
+**Training objective:**
+
+- DINO self-distillation loss with multi-crop
+- iBOT masked-image modeling loss
+- KoLeo regularization on [CLS] tokens
+- Gram anchoring
+
+- **Training regime:** PyTorch FSDP2 (with bf16 and fp8 matrix multiplications)
+
+**Distillation:**
+
+- Distillation follows the standard DINOv3 pretraining procedure, except the teacher is a frozen pretrained ViT-7B.
+
+## Evaluation
+
+**Results**
+
+The reader is referred to the associated paper for details on the evaluation protocols
+
+*Results for ViT backbones pretrained (or distilled) on web (LVD-1689M)*
+
+<table>
+  <tr>
+    <th></th>
+    <!-- <th></th> -->
+    <th colspan="4">Global Tasks</th>
+    <th colspan="5">Dense Tasks</th>
+  </tr>
+  <tr>
+    <th>Model</th>
+    <!-- <th>Dataset</th> -->
+    <th>IN-ReaL</th>
+    <th>IN-R</th>
+    <th>Obj.Net</th>
+    <th>Ox.-H</th>
+    <th>ADE20k</th>
+    <th>NYU↓</th>
+    <th>DAVIS</th>
+    <th>NAVI</th>
+    <th>SPair</th>
+  </tr>
+  <tr>
+    <td>DINOv3 ViT-S/16</td>
+    <!-- <td>LVD-1689M</td> -->
+    <td align="right">87.0</td>
+    <td align="right">60.4</td>
+    <td align="right">50.9</td>
+    <td align="right">49.5</td>
+    <td align="right">47.0</td>
+    <td align="right">0.403</td>
+    <td align="right">72.7</td>
+    <td align="right">56.3</td>
+    <td align="right">50.4</td>
+  </tr>
+  <tr>
+    <td>DINOv3 ViT-S+/16</td>
+    <!-- <td>LVD-1689M</td> -->
+    <td align="right">88.0</td>
+    <td align="right">68.8</td>
+    <td align="right">54.6</td>
+    <td align="right">50.0</td>
+    <td align="right">48.8</td>
+    <td align="right">0.399</td>
+    <td align="right">75.5</td>
+    <td align="right">57.1</td>
+    <td align="right">55.2</td>
+  </tr>
+  <tr>
+    <td>DINOv3 ViT-B/16</td>
+    <!-- <td>LVD-1689M</td> -->
+    <td align="right">89.3</td>
+    <td align="right">76.7</td>
+    <td align="right">64.1</td>
+    <td align="right">58.5</td>
+    <td align="right">51.8</td>
+    <td align="right">0.373</td>
+    <td align="right">77.2</td>
+    <td align="right">58.8</td>
+    <td align="right">57.2</td>
+  </tr>
+  <tr>
+    <td>DINOv3 ViT-L/16</td>
+    <!-- <td>LVD-1689M</td> -->
+    <td align="right">90.2</td>
+    <td align="right">88.1</td>
+    <td align="right">74.8</td>
+    <td align="right">63.1</td>
+    <td align="right">54.9</td>
+    <td align="right">0.352</td>
+    <td align="right">79.9</td>
+    <td align="right">62.3</td>
+    <td align="right">61.3</td>
+  </tr>
+  <tr>
+    <td>DINOv3 ViT-H+/16</td>
+    <!-- <td>LVD-1689M</td> -->
+    <td align="right">90.3</td>
+    <td align="right">90.0</td>
+    <td align="right">78.6</td>
+    <td align="right">64.5</td>
+    <td align="right">54.8</td>
+    <td align="right">0.352</td>
+    <td align="right">79.3</td>
+    <td align="right">63.3</td>
+    <td align="right">56.3</td>
+  </tr>
+  <tr>
+    <td>DINOv3 ViT-7B/16</td>
+    <!-- <td>LVD-1689M</td> -->
+    <td align="right">90.4</td>
+    <td align="right">91.1</td>
+    <td align="right">91.1</td>
+    <td align="right">72.8</td>
+    <td align="right">55.9</td>
+    <td align="right">0.309</td>
+    <td align="right">79.7</td>
+    <td align="right">64.4</td>
+    <td align="right">58.7</td>
+  </tr>
+</table>
+
+*Results for ConvNeXt backbones distilled on web (LVD-1689M)*
+
+<table>
+  <tr>
+    <th></th>
+    <th colspan="6">Global Tasks</th>
+    <th colspan="2">Dense Tasks</th>
+  </tr>
+  <tr>
+    <th>Model</th>
+    <th colspan="2">IN-ReaL</th>
+    <th colspan="2">IN-R</th>
+    <th colspan="2">Obj.Net</th>
+    <th>ADE20k</th>
+    <th>NYU↓</th>
+  </tr>
+  <tr>
+    <td></th>
+    <td>@256px</td>
+    <td>@512px</td>
+    <td>@256px</td>
+    <td>@512px</td>
+    <td>@256px</td>
+    <td>@512px</td>
+    <td colspan="2"></td>
+  </tr>
+  <tr>
+    <td>DINOv3 ConvNeXt Tiny</td>
+    <td align="right">86.6</td>
+    <td align="right">87.7</td>
+    <td align="right">73.7</td>
+    <td align="right">74.1</td>
+    <td align="right">52.6</td>
+    <td align="right">58.7</td>
+    <td align="right">42.7</td>
+    <td align="right">0.448</td>
+  </tr>
+  <tr>
+    <td>DINOv3 ConvNeXt Small</td>
+    <td align="right">87.9</td>
+    <td align="right">88.7</td>
+    <td align="right">73.7</td>
+    <td align="right">74.1</td>
+    <td align="right">52.6</td>
+    <td align="right">58.7</td>
+    <td align="right">44.8</td>
+    <td align="right">0.432</td>
+  </tr>
+  <tr>
+    <td>DINOv3 ConvNeXt Base</td>
+    <td align="right">88.5</td>
+    <td align="right">89.2</td>
+    <td align="right">77.2</td>
+    <td align="right">78.2</td>
+    <td align="right">56.2</td>
+    <td align="right">61.3</td>
+    <td align="right">46.3</td>
+    <td align="right">0.420</td>
+  </tr>
+  <tr>
+    <td>DINOv3 ConvNeXt Large</td>
+    <td align="right">88.9</td>
+    <td align="right">89.4</td>
+    <td align="right">81.3</td>
+    <td align="right">82.4</td>
+    <td align="right">59.3</td>
+    <td align="right">65.2</td>
+    <td align="right">47.8</td>
+    <td align="right">0.403</td>
+  </tr>
+</table>
+
+*Results for ViT backbones pretrained (or distilled) on satellite (SAT-493M)*
+
+<table>
+  <tr>
+    <th></th>
+    <th colspan="7">(GEO-Bench) Classification</th>
+  </tr>
+  <tr>
+    <th>Model</ht>
+    <th>m-BEnet</th>
+    <th>m-brick-kiln
+    <th>m-eurosat</th>
+    <th>m-forestnet</th>
+    <th>m-pv4ger</th>
+    <th>m-so2sat</th>
+    <th>mean</th>
+  </tr>
+  <tr>
+    <td>DINOv3 ViT-L/16</td>
+    <td>73.0</td>
+    <td>96.5</td>
+    <td>94.1</td>
+    <td>60.6</td>
+    <td>96.0</td>
+    <td>57.4</td>
+    <td>79.6</td>
+  </tr>
+  <tr>
+    <td>DINOv3 ViT-7B/16</td>
+    <td>74.0</td>
+    <td>97.2</td>
+    <td>94.8</td>
+    <td>62.3</td>
+    <td>96.1</td>
+    <td>62.1</td>
+    <td>81.1</td>
+  </tr>
+  <tr>
+    <th></th>
+    <th colspan="7">(GEO-Bench) Segmentation</th>
+  </tr>
+  <tr>
+    <th>Model</th>
+    <th>m-cashew</th>
+    <th>m-chesapeake</th>
+    <th>m-NeonTree</th>
+    <th>m-nz-cattle</th>
+    <th>m-pv4ger-seg</th>
+    <th>m-SA-crop</th>
+    <th>mean</th>
+  </tr>
+  <tr>
+    <td>DINOv3 ViT-L/16</td>
+    <td>94.2</td>
+    <td>75.6</td>
+    <td>61.8</td>
+    <td>83.7</td>
+    <td>95.2</td>
+    <td>36.8</td>
+    <td>74.5</td>
+  </tr>
+  <tr>
+    <td>DINOv3 ViT-7B/16</td>
+    <td>94.1</td>
+    <td>76.6</td>
+    <td>62.6</td>
+    <td>83.4</td>
+    <td>95.5</td>
+    <td>37.6</td>
+    <td>75.0</td>
+  </tr>
+</table>
+
+
+## Environmental Impact
+
+- **Hardware Type:** Nvidia H100
+- **Hours used:** 61,440 hours for ViT-7B model training
+- **Cloud Provider:** Private infrastructure
+- **Compute Region:** USA
+- **Carbon Emitted:** 18t CO2eq
+
+## Technical Specifications
+
+### Model Architecture and Objective
+
+Vision Transformer models:
+
+- ViT-S (21M parameters): patch size 16, embedding dimension 384, 4 register tokens, 6 heads, MLP FFN, RoPE
+- ViT-S+ (29M parameters): patch size 16, embedding dimension 384, 4 register tokens, 6 heads, SwiGLU FFN, RoPE
+- ViT-B (86M parameters): patch size 16, embedding dimension 768, 4 register tokens, 12 heads, MLP FFN, RoPE
+- ViT-L (300M parameters): patch size 16, embedding dimension 1024, 4 register tokens, 16 heads, MLP FFN, RoPE
+- ViT-H+ (840M parameters): patch size 16, embedding dimension 1280, 4 register tokens, 20 heads, SwiGLU FFN, RoPE
+- ViT-7B (6716M parameters): patch size 16, embedding dimension 4096, 4 register tokens, 32 heads, SwiGLU FFN, RoPE
+
+ConvNeXt models:
+
+- ConvNeXt Tiny (29M parameters)
+- ConvNeXt Small (50M parameters)
+- ConvNeXt Base (89M parameters)
+- ConvNeXt Large (198M parameters)
+
+### Compute Infrastructure
+
+#### Hardware
+
+Nvidia H100 GPUs
+
+#### Software
+
+PyTorch 2.7
+
+## More Information
+
+See the [blog post](https://ai.meta.com/blog/dinov3-self-supervised-vision-model/) and the associated [website](https://ai.meta.com/dinov3/).
+
+## Citation
+
+**BibTeX**
+
+```
+@misc{simeoni2025dinov3,
+  title={{DINOv3}},
+  author={Sim{\'e}oni, Oriane and Vo, Huy V. and Seitzer, Maximilian and Baldassarre, Federico and Oquab, Maxime and Jose, Cijo and Khalidov, Vasil and Szafraniec, Marc and Yi, Seungeun and Ramamonjisoa, Micha{\"e}l and Massa, Francisco and Haziza, Daniel and Wehrstedt, Luca and Wang, Jianyuan and Darcet, Timoth{\'e}e and Moutakanni, Th{\'e}o and Sentana, Leonel and Roberts, Claire and Vedaldi, Andrea and Tolan, Jamie and Brandt, John and Couprie, Camille and Mairal, Julien and J{\'e}gou, Herv{\'e} and Labatut, Patrick and Bojanowski, Piotr},
+  year={2025},
+  eprint={2508.10104},
+  archivePrefix={arXiv},
+  primaryClass={cs.CV},
+  url={https://arxiv.org/abs/2508.10104},
+}
+```

InfiniDepth/model/block/torchhub/dinov3/README.md

@@ -0,0 +1,734 @@
+🆕 [2025-08-14] :fire: DINOv3 backbones are now available in [Hugging Face Hub](https://huggingface.co/collections/facebook/dinov3-68924841bd6b561778e31009) and [supported](https://huggingface.co/docs/transformers/model_doc/dinov3) by the Hugging Face [Transformers](https://huggingface.co/docs/transformers/index) library
+
+# DINOv3 🦖🦖🦖
+
+**[Meta AI Research, FAIR](https://ai.meta.com/research/)**
+
+Oriane Siméoni, Huy V. Vo, Maximilian Seitzer, Federico Baldassarre, Maxime Oquab, <br/>
+Cijo Jose, Vasil Khalidov, Marc Szafraniec, Seungeun Yi, Michaël Ramamonjisoa, <br/>
+Francisco Massa, Daniel Haziza, Luca Wehrstedt, Jianyuan Wang, <br/>
+Timothée Darcet, Théo Moutakanni, Leonel Sentana, Claire Roberts, <br/>
+Andrea Vedaldi, Jamie Tolan, John Brandt, Camille Couprie, <br/>
+Julien Mairal, Hervé Jégou, Patrick Labatut, Piotr Bojanowski
+
+[ :scroll: [`Paper`](https://arxiv.org/abs/2508.10104)] [ :newspaper: [`Blog`](https://ai.meta.com/blog/dinov3-self-supervised-vision-model/)] [ :globe_with_meridians: [`Website`](https://ai.meta.com/dinov3/)] [ :book: [`BibTeX`](#citing-dinov3)]
+
+Reference PyTorch implementation and models for DINOv3. For details, see the **[DINOv3](https://arxiv.org/abs/2508.10104)** paper.
+
+## Overview
+
+<div align="center">
+  <img width="1364" height="1024" alt="market" src="https://github.com/user-attachments/assets/1411f491-988e-49cb-95ae-d03fe6e3c268" />
+
+  <i></em><b>High-resolution dense features.</b><br/>We visualize the cosine similarity maps obtained with DINOv3 output features<br/> between the patches marked with a red cross and all other patches.</i>
+</div>
+
+<br/>
+
+An extended family of versatile vision foundation models producing high-quality dense features and achieving outstanding performance on various vision tasks including outperforming the specialized state of the art across a broad range of settings, without fine-tuning
+
+## Pretrained models
+
+:information_source: Please follow the link provided below to get access to all the model weights: once accepted, an e-mail will be sent with the complete list of URLs pointing to all the available model weights (both backbones and adapters). These URLs can then be used to either:
+- download the model or adapter weights to a local filesystem and point `torch.hub.load()` to these local weights via the `weights` or `backbone_weights` parameters, or
+- directly invoke `torch.hub.load()` to download and load a backbone or an adapter from its URL via also the `weights` or `backbone_weights` parameters.
+
+See the example code snippets below.
+
+:warning: Please use `wget` instead of a web browser to download the weights.
+
+ViT models pretrained on web dataset (LVD-1689M):
+<table style="margin: auto">
+  <thead>
+    <tr>
+      <th>Model</th>
+      <th>Parameters</th>
+      <th>Pretraining<br/>Dataset</th>
+      <th>Download</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>ViT-S/16 distilled </td>
+      <td align="right">21M</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+    <tr>
+      <td>ViT-S+/16 distilled</td>
+      <td align="right">29M</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+    <tr>
+      <td>ViT-B/16 distilled</td>
+      <td align="right">86M</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+    <tr>
+      <td>ViT-L/16 distilled</td>
+      <td align="right">300M</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+    <tr>
+      <td>ViT-H+/16 distilled</td>
+      <td align="right">840M</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+    <tr>
+      <td>ViT-7B/16</td>
+      <td align="right">6,716M</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+  </tbody>
+</table>
+
+ConvNeXt models pretrained on web dataset (LVD-1689M):
+<table style="margin: auto">
+  <thead>
+    <tr>
+      <th>Model</th>
+      <th>Parameters</th>
+      <th>Pretraining<br/>Dataset</th>
+      <th>Download</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>ConvNeXt Tiny</td>
+      <td align="right">29M</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+    <tr>
+      <td>ConvNeXt Small</td>
+      <td align="right">50M</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+    <tr>
+      <td>ConvNeXt Base</td>
+      <td align="right">89M</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+    <tr>
+      <td>ConvNeXt Large</td>
+      <td align="right">198M</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+  </tbody>
+</table>
+
+ViT models pretrained on satellite dataset (SAT-493M):
+<table style="margin: auto">
+  <thead>
+    <tr>
+      <th>Model</th>
+      <th>Parameters</th>
+      <th>Pretraining<br/>Dataset</th>
+      <th>Download</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>ViT-L/16 distilled</td>
+      <td align="right">300M</td>
+      <td align="center">SAT-493M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+    <tr>
+      <td>ViT-7B/16</td>
+      <td align="right">6,716M</td>
+      <td align="center">SAT-493M</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+  </tbody>
+</table>
+
+
+### Pretrained backbones (via PyTorch [Hub](https://docs.pytorch.org/docs/stable/hub.html))
+
+Please follow the instructions [here](https://pytorch.org/get-started/locally/) to install PyTorch (the only required dependency for loading the model). Installing PyTorch with CUDA support is strongly recommended.
+
+```python
+import torch
+
+REPO_DIR = <PATH/TO/A/LOCAL/DIRECTORY/WHERE/THE/DINOV3/REPO/WAS/CLONED>
+
+# DINOv3 ViT models pretrained on web images
+dinov3_vits16 = torch.hub.load(REPO_DIR, 'dinov3_vits16', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+dinov3_vits16plus = torch.hub.load(REPO_DIR, 'dinov3_vits16plus', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+dinov3_vitb16 = torch.hub.load(REPO_DIR, 'dinov3_vitb16', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+dinov3_vitl16 = torch.hub.load(REPO_DIR, 'dinov3_vitl16', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+dinov3_vith16plus = torch.hub.load(REPO_DIR, 'dinov3_vith16plus', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+dinov3_vit7b16 = torch.hub.load(REPO_DIR, 'dinov3_vit7b16', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+
+# DINOv3 ConvNeXt models pretrained on web images
+dinov3_convnext_tiny = torch.hub.load(REPO_DIR, 'dinov3_convnext_tiny', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+dinov3_convnext_small = torch.hub.load(REPO_DIR, 'dinov3_convnext_small', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+dinov3_convnext_base = torch.hub.load(REPO_DIR, 'dinov3_convnext_base', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+dinov3_convnext_large = torch.hub.load(REPO_DIR, 'dinov3_convnext_large', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+
+# DINOv3 ViT models pretrained on satellite imagery
+dinov3_vitl16 = torch.hub.load(REPO_DIR, 'dinov3_vitl16', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+dinov3_vit7b16 = torch.hub.load(REPO_DIR, 'dinov3_vit7b16', source='local', weights=<CHECKPOINT/URL/OR/PATH>)
+```
+
+### Pretrained backbones (via Hugging Face [Transformers](https://huggingface.co/docs/transformers/))
+
+All the backbones are available in the the [DINOv3](https://huggingface.co/collections/facebook/dinov3-68924841bd6b561778e31009) collection on Hugging Face Hub and supported via the Hugging Face [Transformers](https://huggingface.co/docs/transformers/index) library. Please refer to the corresponding documentation for usage, but below is a short example that demonstrates how to obtain an image embedding with either [Pipeline] or the [AutoModel] class.
+
+```python
+from transformers import pipeline
+from transformers.image_utils import load_image
+
+url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
+image = load_image(url)
+
+feature_extractor = pipeline(
+    model="facebook/dinov3-convnext-tiny-pretrain-lvd1689m",
+    task="image-feature-extraction", 
+)
+features = feature_extractor(image)
+```
+
+```python
+import torch
+from transformers import AutoImageProcessor, AutoModel
+from transformers.image_utils import load_image
+
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = load_image(url)
+
+pretrained_model_name = "facebook/dinov3-convnext-tiny-pretrain-lvd1689m"
+processor = AutoImageProcessor.from_pretrained(pretrained_model_name)
+model = AutoModel.from_pretrained(
+    pretrained_model_name, 
+    device_map="auto", 
+)
+
+inputs = processor(images=image, return_tensors="pt").to(model.device)
+with torch.inference_mode():
+    outputs = model(**inputs)
+
+pooled_output = outputs.pooler_output
+print("Pooled output shape:", pooled_output.shape)
+```
+
+where `model` and `pretrained_model_name` above can be one of:
+- `facebook/dinov3-vits16-pretrain-lvd1689m`
+- `facebook/dinov3-vits16plus-pretrain-lvd1689m`
+- `facebook/dinov3-vitb16-pretrain-lvd1689m`
+- `facebook/dinov3-vitl16-pretrain-lvd1689m`
+- `facebook/dinov3-vith16plus-pretrain-lvd1689m`
+- `facebook/dinov3-vit7b16-pretrain-lvd1689m`
+- `facebook/dinov3-convnext-base-pretrain-lvd1689m`
+- `facebook/dinov3-convnext-large-pretrain-lvd1689m`
+- `facebook/dinov3-convnext-small-pretrain-lvd1689m`
+- `facebook/dinov3-convnext-tiny-pretrain-lvd1689m`
+- `facebook/dinov3-vitl16-pretrain-sat493m`
+- `facebook/dinov3-vit7b16-pretrain-sat493m`
+
+### Image transforms
+
+For models using the LVD-1689M weights (pretrained on web images), please use the following transform (standard ImageNet evaluation transform):
+
+```python
+import torchvision
+
+def make_transform(resize_size: int = 224):
+    to_tensor = transforms.ToTensor()
+    resize = transforms.Resize((resize_size, resize_size), antialias=True)
+    normalize = transforms.Normalize(
+        mean=(0.485, 0.456, 0.406),
+        std=(0.229, 0.224, 0.225),
+    )
+    return transforms.Compose([to_tensor, resize, normalize])
+```
+
+
+For models using the SAT-493M weights (pretrained on satellite imagery), please use the following transform:
+
+
+```python
+import torchvision
+
+def make_transform(resize_size: int = 224):
+    to_tensor = transforms.ToTensor()
+    resize = transforms.Resize((resize_size, resize_size), antialias=True)
+    normalize = transforms.Normalize(
+        mean=(0.430, 0.411, 0.296),
+        std=(0.213, 0.156, 0.143),
+    )
+    return transforms.Compose([to_tensor, resize, normalize])
+```
+
+### Pretrained heads - Image classification
+
+<table style="margin: auto">
+  <thead>
+    <tr>
+      <th>Backbone</th>
+      <th>Pretraining<br/>Dataset</th>
+      <th>Head<br/>Dataset</th>
+      <th>Download</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>ViT-7B/16</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center">ImageNet</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+  </tbody>
+</table>
+
+
+The (full) classifier models can be loaded via PyTorch Hub:
+
+```python
+import torch
+
+# DINOv3
+dinov3_vit7b16_lc = torch.hub.load(REPO_DIR, 'dinov3_vit7b16_lc', source="local", weights=<DEPTHER/CHECKPOINT/URL/OR/PATH>, backbone_weights=<BACKBONE/CHECKPOINT/URL/OR/PATH>)
+
+```
+
+### Pretrained heads - Depther trained on SYNTHMIX dataset
+
+<table style="margin: auto">
+  <thead>
+    <tr>
+      <th>Backbone</th>
+      <th>Pretraining<br/>Dataset</th>
+      <th>Head<br/>Dataset</th>
+      <th>Download</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>ViT-7B/16</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center">SYNTHMIX</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+  </tbody>
+</table>
+
+
+```python
+depther = torch.hub.load(REPO_DIR, 'dinov3_vit7b16_dd', source="local", weights=<DEPTHER/CHECKPOINT/URL/OR/PATH>, backbone_weights=<BACKBONE/CHECKPOINT/URL/OR/PATH>)
+```
+
+Full example code of depther on an image
+
+```python
+from PIL import Image
+import torch
+from torchvision import transforms
+import matplotlib.pyplot as plt
+from matplotlib import colormaps
+
+def get_img():
+    import requests
+    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+    image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
+    return image
+
+def make_transform(resize_size: int | list[int] = 768):
+    to_tensor = transforms.ToTensor()
+    resize = transforms.Resize((resize_size, resize_size), antialias=True)
+    normalize = transforms.Normalize(
+        mean=(0.485, 0.456, 0.406),
+        std=(0.229, 0.224, 0.225),
+    )
+    return transforms.Compose([to_tensor, resize, normalize])
+
+depther = torch.hub.load(REPO_DIR, 'dinov3_vit7b16_dd', source="local", weights=<DEPTHER/CHECKPOINT/URL/OR/PATH>, backbone_weights=<BACKBONE/CHECKPOINT/URL/OR/PATH>)
+
+img_size = 1024
+img = get_img()
+transform = make_transform(img_size)
+with torch.inference_mode():
+    with torch.autocast('cuda', dtype=torch.bfloat16):
+        batch_img = transform(img)[None]
+        batch_img = batch_img
+        depths = depther(batch_img)
+
+plt.figure(figsize=(12, 6))
+plt.subplot(121)
+plt.imshow(img)
+plt.axis("off")
+plt.subplot(122)
+plt.imshow(depths[0,0].cpu(), cmap=colormaps["Spectral"])
+plt.axis("off")
+
+```
+
+### Pretrained heads - Detector trained on COCO2017 dataset
+
+<table style="margin: auto">
+  <thead>
+    <tr>
+      <th>Backbone</th>
+      <th>Pretraining<br/>Dataset</th>
+      <th>Head<br/>Dataset</th>
+      <th>Download</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>ViT-7B/16</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center">COCO2017</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+  </tbody>
+</table>
+
+
+```python
+detector = torch.hub.load(REPO_DIR, 'dinov3_vit7b16_de', source="local", weights=<DETECTOR/CHECKPOINT/URL/OR/PATH>, backbone_weights=<BACKBONE/CHECKPOINT/URL/OR/PATH>)
+```
+
+### Pretrained heads - Segmentor trained on ADE20K dataset
+
+<table style="margin: auto">
+  <thead>
+    <tr>
+      <th>Backbone</th>
+      <th>Pretraining<br/>Dataset</th>
+      <th>Head<br/>Dataset</th>
+      <th>Download</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>ViT-7B/16</td>
+      <td align="center">LVD-1689M</td>
+      <td align="center">ADE20K</td>
+      <td align="center"><a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a></td>
+    </tr>
+  </tbody>
+</table>
+
+```python
+segmentor = torch.hub.load(REPO_DIR, 'dinov3_vit7b16_ms', source="local", weights=<SEGMENTOR/CHECKPOINT/URL/OR/PATH>, backbone_weights=<BACKBONE/CHECKPOINT/URL/OR/PATH>)
+```
+
+Full example code of segmentator on an image
+
+```python
+import sys
+sys.path.append(REPO_DIR)
+
+from PIL import Image
+import torch
+from torchvision import transforms
+import matplotlib.pyplot as plt
+from matplotlib import colormaps
+from functools import partial
+from dinov3.eval.segmentation.inference import make_inference
+
+
+def get_img():
+    import requests
+    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+    image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
+    return image
+
+def make_transform(resize_size: int | list[int] = 768):
+    to_tensor = transforms.ToTensor()
+    resize = transforms.Resize((resize_size, resize_size), antialias=True)
+    normalize = transforms.Normalize(
+        mean=(0.485, 0.456, 0.406),
+        std=(0.229, 0.224, 0.225),
+    )
+    return transforms.Compose([to_tensor, resize, normalize])
+
+segmentor = torch.hub.load(REPO_DIR, 'dinov3_vit7b16_ms', source="local", weights=<SEGMENTOR/CHECKPOINT/URL/OR/PATH>, backbone_weights=<BACKBONE/CHECKPOINT/URL/OR/PATH>)
+
+img_size = 896
+img  = get_img()
+transform = make_transform(img_size)
+with torch.inference_mode():
+    with torch.autocast('cuda', dtype=torch.bfloat16):
+        batch_img = transform(img)[None]
+        pred_vit7b = segmentor(batch_img)  # raw predictions  
+        # actual segmentation map
+        segmentation_map_vit7b = make_inference(
+            batch_img,
+            segmentor,
+            inference_mode="slide",
+            decoder_head_type="m2f",
+            rescale_to=(img.size[-1], img.size[-2]),
+            n_output_channels=150,
+            crop_size=(img_size, img_size),
+            stride=(img_size, img_size),
+            output_activation=partial(torch.nn.functional.softmax, dim=1),
+        ).argmax(dim=1, keepdim=True)
+plt.figure(figsize=(12, 6))
+plt.subplot(121)
+plt.imshow(img)
+plt.axis("off")
+plt.subplot(122)
+plt.imshow(segmentation_map_vit7b[0,0].cpu(), cmap=colormaps["Spectral"])
+plt.axis("off")
+```
+
+
+
+
+### Pretrained heads - Zero-shot tasks with `dino.txt`
+
+<table style="margin: auto">
+  <thead>
+    <tr>
+      <th rowspan="2">Backbone</th>
+      <th>Download</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td>ViT-L/16 distilled</td>
+      <td align="center">
+        <a href="https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/">[link]</a>,
+        <a href="https://dl.fbaipublicfiles.com/dinov3/thirdparty/bpe_simple_vocab_16e6.txt.gz">vocabulary</a>,
+        <a href="https://dl.fbaipublicfiles.com/dinov2/thirdparty/LICENSE">vocabulary license</a>
+      </td>
+    </tr>
+  </tbody>
+</table>
+
+The (full) dino.txt model can be loaded via PyTorch Hub:
+
+```python
+import torch
+# DINOv3
+dinov3_vitl16_dinotxt_tet1280d20h24l, tokenizer = torch.hub.load(REPO_DIR, 'dinov3_vitl16_dinotxt_tet1280d20h24l', weights=<SEGMENTOR/CHECKPOINT/URL/OR/PATH>, backbone_weights=<BACKBONE/CHECKPOINT/URL/OR/PATH>)
+```
+
+
+## Installation
+
+The training and evaluation code requires PyTorch version >= 2.7.1 as well as a few other 3rd party packages. Note that the code has only been tested with the specified versions and also expects a Linux environment. To setup all the required dependencies for training and evaluation, please follow the instructions below:
+
+*[micromamba](https://mamba.readthedocs.io/en/latest/user_guide/micromamba.html)* **(Recommended)** - Clone the repository and then create and activate a `dinov3` conda environment using the provided environment definition:
+
+```shell
+micromamba env create -f conda.yaml
+micromamba activate dinov3
+```
+
+## Getting started
+
+Several notebooks are provided to get started applying DINOv3:
+- [PCA of patch features](notebooks/pca.ipynb): display the PCA of DINOv3 patch features on a foreground object (rainbow visualizations from the paper) [[Run in Google Colab]](https://colab.research.google.com/github/facebookresearch/dinov3/blob/main/notebooks/pca.ipynb)
+- [Foreground segmentation](notebooks/foreground_segmentation.ipynb): train a linear foreground segmentation model based on DINOv3 features [[Run in Google Colab]](https://colab.research.google.com/github/facebookresearch/dinov3/blob/main/notebooks/foreground_segmentation.ipynb)
+- [Dense and sparse matching](notebooks/dense_sparse_matching.ipynb): match patches from objects on two different images based on DINOv3 features [[Run in Google Colab]](https://colab.research.google.com/github/facebookresearch/dinov3/blob/main/notebooks/dense_sparse_matching.ipynb)
+- [Segmentation tracking](notebooks/segmentation_tracking.ipynb): video segmentation tracking using a non-parametric method based on DINOv3 features [[Run in Google Colab]](https://colab.research.google.com/github/facebookresearch/dinov3/blob/main/notebooks/segmentation_tracking.ipynb)
+
+## Data preparation
+
+### ImageNet-1k
+
+The root directory of the dataset should hold the following contents:
+
+- `<ROOT>/test/ILSVRC2012_test_00000001.JPEG`
+- `<ROOT>/test/[..]`
+- `<ROOT>/test/ILSVRC2012_test_00100000.JPEG`
+- `<ROOT>/train/n01440764/n01440764_10026.JPEG`
+- `<ROOT>/train/[...]`
+- `<ROOT>/train/n15075141/n15075141_9993.JPEG`
+- `<ROOT>/val/n01440764/ILSVRC2012_val_00000293.JPEG`
+- `<ROOT>/val/[...]`
+- `<ROOT>/val/n15075141/ILSVRC2012_val_00049174.JPEG`
+- `<ROOT>/labels.txt`
+
+The provided dataset implementation expects a few additional metadata files to be present under the extra directory:
+
+- `<EXTRA>/class-ids-TRAIN.npy`
+- `<EXTRA>/class-ids-VAL.npy`
+- `<EXTRA>/class-names-TRAIN.npy`
+- `<EXTRA>/class-names-VAL.npy`
+- `<EXTRA>/entries-TEST.npy`
+- `<EXTRA>/entries-TRAIN.npy`
+- `<EXTRA>/entries-VAL.npy`
+
+These metadata files can be generated (once) with the following lines of Python code:
+
+```python
+from dinov3.data.datasets import ImageNet
+
+for split in ImageNet.Split:
+    dataset = ImageNet(split=split, root="<ROOT>", extra="<EXTRA>")
+    dataset.dump_extra()
+```
+
+Note that the root and extra directories do not have to be distinct directories.
+
+### ImageNet-22k
+
+Please adapt the [dataset class](dinov3/data/datasets/image_net_22k.py) to match your local setup.
+
+<br />
+
+:warning: To execute the commands provided in the next sections for training and evaluation, the `dinov3` package should be included in the Python module search path, i.e. simply prefix the command to run with `PYTHONPATH=.`.
+
+## Training
+
+### Fast setup: training DINOv3 ViT-L/16 on ImageNet-1k
+
+Run DINOv3 pre-training on 4 H100-80GB nodes (32 GPUs) in a SLURM cluster environment with submitit:
+
+```shell
+ PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/train/train.py \
+  --nodes 4 \
+  --config-file dinov3/configs/train/vitl_im1k_lin834.yaml \
+  --output-dir <PATH/TO/OUTPUT/DIR> \
+  train.dataset_path=ImageNet22k:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
+```
+Training time is approximately 14 hours and the resulting checkpoint should reach 82.0% on k-NN eval and 83.5% on linear eval.
+
+The training code saves the weights of the teacher in the eval folder every 12500 iterations for evaluation.
+
+### Exact DINOv3 setup: training DINOv3 ViT-7B/16
+
+DINOv3 ViT-7B/16 is trained on a private dataset. The training involves 3 stages:
+- Pretraining
+- Gram anchoring
+- High resolution adaptation
+
+#### Pretraining
+
+Launch DINOV3 ViT-7B/16 pretraining on 32 nodes (256 GPUs) in a SLURM cluster environment with submitit.
+
+```shell
+PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/train/train.py \
+  --nodes 32 \
+  --config-file dinov3/configs/train/dinov3_vit7b16_pretrain.yaml \
+  --output-dir <PATH/TO/OUTPUT/DIR> \
+  train.dataset_path=<DATASET>:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
+```
+
+#### Gram anchoring
+
+```shell
+PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/train/train.py \
+  --nodes 32 \
+  --config-file dinov3/configs/train/dinov3_vit7b16_gram_anchor.yaml \
+  --output-dir <PATH/TO/OUTPUT/DIR> \
+  train.dataset_path=<DATASET>:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET> \
+  gram.ckpt=<PATH/TO/GRAM_TEACHER_FROM_PREVIOUS_STEP>   
+```
+
+#### High-resolution adaptation
+
+
+```shell
+PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/train/train.py \
+  --nodes 32 \
+  --config-file dinov3/configs/train/dinov3_vit7b16_high_res_adapt.yaml \
+  --output-dir <PATH/TO/OUTPUT/DIR> \
+  train.dataset_path=<DATASET>:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET> \
+  gram.ckpt=<PATH/TO/TEACHER_FROM_GRAM> \
+  student.resume_from_teacher_chkpt=<PATH/TO/TEACHER_FROM_GRAM>
+```
+
+## Multi-distillation 
+
+### Test setup:
+
+```shell
+PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/train/train.py \
+  --nodes 1 \
+  --config-file dinov3/configs/train/multi_distillation_test.yaml \
+  --output-dir <PATH/TO/OUTPUT/DIR> \
+  --multi-distillation \
+  train.dataset_path=<DATASET>:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
+```
+
+## Evaluation
+
+The training code regularly saves the teacher weights. In order to evaluate the model, run the following evaluation on a single node:
+
+
+### Logistic regression classification on ImageNet-1k
+
+```shell
+PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/eval/log_regression.py \
+  model.config_file=<PATH/TO/OUTPUT/DIR>/config.yaml \
+  model.pretrained_weights=<PATH/TO/OUTPUT/DIR>/teacher_checkpoint.pth \
+  output_dir=<PATH/TO/OUTPUT/DIR> \
+  train.dataset=ImageNet:split=TRAIN:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET> \
+  eval.test_dataset=ImageNet:split=VAL:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
+```
+
+### k-NN classification on ImageNet-1k
+
+```shell
+PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/eval/knn.py \
+  model.config_file=<PATH/TO/OUTPUT/DIR>/config.yaml \
+  model.pretrained_weights=<PATH/TO/OUTPUT/DIR>/teacher_checkpoint.pth \
+  output_dir=<PATH/TO/OUTPUT/DIR> \
+  train.dataset=ImageNet:split=TRAIN:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET> \
+  eval.test_dataset=ImageNet:split=VAL:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
+```
+
+### Linear classification with data augmentation on ImageNet-1k
+
+```shell
+PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/eval/linear.py \
+  model.config_file=<PATH/TO/OUTPUT/DIR>/config.yaml \
+  model.pretrained_weights=<PATH/TO/OUTPUT/DIR>/teacher_checkpoint.pth \
+  output_dir=<PATH/TO/OUTPUT/DIR> \
+  train.dataset=ImageNet:split=TRAIN:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET> \
+  train.val_dataset=ImageNet:split=VAL:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
+```
+
+
+### Text alignment on DINOv3 using dino.txt
+
+Text alignment can be done following the method from `dino.txt` aka [DINOv2 Meets Text](https://arxiv.org/abs/2412.16334).
+
+```shell
+PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/eval/text/train_dinotxt.py \
+   --nodes 4 \
+  # An example config for text alignment is here: dinov3/eval/text/configs/dinov3_vitl_text.yaml \ 
+  trainer_config_file="<PATH/TO/DINOv3/TEXT/CONFIG>" \
+  output-dir=<PATH/TO/OUTPUT/DIR>
+```
+Launching the above trains text alignment on 4 nodes with 8 gpus each (32 gpus in total).
+Please note that the text alignment model in the DINOv3 paper was trained on a private dataset and here we have given an example config in ```dinov3/eval/text/configs/dinov3_vitl_text.yaml``` using ```CocoCaptions``` dataset for illustration purposes.
+Please adapt the provided ```CocoCaptions``` dataset class, the dataset can be found [here](https://www.kaggle.com/datasets/nikhil7280/coco-image-caption)  
+
+## License
+
+DINOv3 code and model weights are released under the DINOv3 License. See [LICENSE.md](LICENSE.md) for additional details.
+
+## Contributing
+
+See [contributing](CONTRIBUTING.md) and the [code of conduct](CODE_OF_CONDUCT.md).
+
+## Citing DINOv3
+
+If you find this repository useful, please consider giving a star :star: and citation :t-rex::
+
+```
+@misc{simeoni2025dinov3,
+  title={{DINOv3}},
+  author={Sim{\'e}oni, Oriane and Vo, Huy V. and Seitzer, Maximilian and Baldassarre, Federico and Oquab, Maxime and Jose, Cijo and Khalidov, Vasil and Szafraniec, Marc and Yi, Seungeun and Ramamonjisoa, Micha{\"e}l and Massa, Francisco and Haziza, Daniel and Wehrstedt, Luca and Wang, Jianyuan and Darcet, Timoth{\'e}e and Moutakanni, Th{\'e}o and Sentana, Leonel and Roberts, Claire and Vedaldi, Andrea and Tolan, Jamie and Brandt, John and Couprie, Camille and Mairal, Julien and J{\'e}gou, Herv{\'e} and Labatut, Patrick and Bojanowski, Piotr},
+  year={2025},
+  eprint={2508.10104},
+  archivePrefix={arXiv},
+  primaryClass={cs.CV},
+  url={https://arxiv.org/abs/2508.10104},
+}
+```

InfiniDepth/model/block/torchhub/dinov3/conda.yaml

@@ -0,0 +1,23 @@
+name: dinov3
+channels:
+  - defaults
+  - conda-forge
+dependencies:
+  - python=3.11
+  - omegaconf
+  - pip
+  - pip:
+    - ftfy # needed for dino.txt
+    - iopath
+    - omegaconf
+    - pandas 
+    - regex # needed for dino.txt
+    - pandas
+    - scikit-learn
+    - scikit-learn-intelex
+    - submitit
+    - termcolor
+    - torch 
+    - torchvision 
+    - torchmetrics
+    

InfiniDepth/model/block/torchhub/dinov3/dinov3/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+__version__ = "0.0.1"

InfiniDepth/model/block/torchhub/dinov3/dinov3/checkpointer/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from .checkpointer import (
+    CheckpointRetentionPolicy,
+    cleanup_checkpoint,
+    find_all_checkpoints,
+    find_latest_checkpoint,
+    init_fsdp_model_from_checkpoint,
+    init_model_from_checkpoint_for_evals,
+    keep_checkpoint_copy,
+    keep_last_n_checkpoints,
+    load_checkpoint,
+    register_dont_save_hooks,
+    save_checkpoint,
+)

InfiniDepth/model/block/torchhub/dinov3/dinov3/checkpointer/checkpointer.py

@@ -0,0 +1,349 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+"""
+Suggested file structure:
+
+output_dir/
+|-- ckpt/
+|   |-- 0/
+|   |-- 99/
+|   |-- 199/
+|   |-- 199_keep/
+|   |-- 299/
+|   `-- ...
+`-- eval/
+    `-- 0/
+    `-- 99/
+        `-- ckpt/
+
+Distributed checkpointer docs:
+- https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html
+- https://pytorch.org/docs/stable/distributed.checkpoint.html
+"""
+
+import logging
+import shutil
+import subprocess
+import tempfile
+from enum import Enum
+from pathlib import Path
+from typing import List, Sequence, Set
+
+import torch
+import torch.distributed as dist
+import torch.distributed.checkpoint as dcp
+import torch.distributed.checkpoint.filesystem as dcpfs
+import torch.distributed.checkpoint.state_dict as dcpsd
+from torch.distributed.checkpoint.stateful import Stateful
+
+logger = logging.getLogger(__name__)
+
+
+class CheckpointRetentionPolicy(Enum):
+    ALL = "all"  # keep all checkpoints
+    BEST = "best"
+    LAST = "last"
+    LAST_AND_BEST = "last_and_best"
+    NONE = "none"  # do not keep any checkpoints
+
+    @property
+    def keep_filters(self) -> Set[str]:
+        """Files that match these patterns are not deleted by cleanup"""
+        if self == CheckpointRetentionPolicy.LAST:
+            return set(["final"])
+        if self == CheckpointRetentionPolicy.BEST:
+            return set(["best"])
+        if self == CheckpointRetentionPolicy.LAST_AND_BEST:
+            return set(["final", "best"])
+        if self == CheckpointRetentionPolicy.ALL:
+            return set()
+        return set()
+
+    @property
+    def max_to_keep(self) -> int | None:
+        """
+        maximum "periodic" checkpoints to keep concurrently, ie. saved with `step` and not `save`. `None` for keep all
+        """
+        if self == CheckpointRetentionPolicy.ALL:
+            return None
+        return 1
+
+
+def save_checkpoint(
+    ckpt_dir: str | Path,  # output_dir/ckpt/199
+    *,
+    iteration: int | str,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer | None = None,
+    overwrite: bool = True,
+    process_group: dist.ProcessGroup = None,
+    **others: Stateful,
+):
+    """Save a plain/DDP/FSDP/FSDP2 model, its optimizer, an integer iteration and other stateful objects."""
+    rank = torch.distributed.get_rank(group=process_group)
+
+    # Rank 0 checks if the checkpoint directory exists, but all ranks need to know if if exists,
+    # so they can raise an error when overwrite is False. If overwrite is True, rank 0 will delete it
+    # and other ranks wait for the deletion to finish.
+    ckpt_dir = Path(ckpt_dir)
+    ckpt_dir_exists = [ckpt_dir.exists() if rank == 0 else None]
+    src_rank = 0
+    if process_group is not None:
+        src_rank = torch.distributed.get_global_rank(group=process_group, group_rank=0)
+    torch.distributed.broadcast_object_list(ckpt_dir_exists, src=src_rank, group=process_group)
+    ckpt_dir_exists = ckpt_dir_exists[0]
+    if ckpt_dir_exists:
+        if overwrite:
+            if rank == 0:
+                if ckpt_dir.is_dir():
+                    shutil.rmtree(ckpt_dir)
+                else:
+                    ckpt_dir.unlink()
+                logger.info(f"Deleted: {ckpt_dir}")
+            torch.distributed.barrier(group=process_group)
+        else:
+            raise RuntimeError(f"Checkpoint already exists: {ckpt_dir}")
+
+    # Rank 0 creates a temporary directory for the checkpoint and broadcasts the name to all ranks.
+    ckpt_dir.parent.mkdir(parents=True, exist_ok=True)
+    ckpt_dir_tmp = [tempfile.mkdtemp(dir=ckpt_dir.parent, prefix=ckpt_dir.name) if rank == 0 else None]
+    torch.distributed.broadcast_object_list(ckpt_dir_tmp, src=src_rank, group=process_group)
+    ckpt_dir_tmp = Path(ckpt_dir_tmp[0])
+
+    to_save = {"iteration": iteration}
+    to_save["model"] = dcpsd.get_model_state_dict(model)
+    if optimizer is not None:
+        to_save["optimizer"] = dcpsd.get_optimizer_state_dict(model, optimizer)
+    to_save.update(others)
+    dcp.save(
+        to_save,
+        storage_writer=dcpfs.FileSystemWriter(ckpt_dir_tmp),
+        process_group=process_group,
+    )
+
+    # Rank 0 renames the temporary directory to the final checkpoint directory. All ranks wait for the rename.
+    if rank == 0:
+        ckpt_dir_tmp.rename(ckpt_dir)
+    torch.distributed.barrier()
+
+    logger.info(f"Saved: {ckpt_dir}")
+
+
+def load_checkpoint(
+    ckpt_dir: str | Path,  # output_dir/ckpt/199
+    *,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer | None = None,
+    strict_loading: bool = True,
+    process_group: dist.ProcessGroup = None,
+    **others: Stateful,
+) -> int | None:
+    """
+    Load a plain/DDP/FSDP/FSDP2 model, its optimizer, an integer iteration and other stateful objects.
+    Can you take a checkpoint saved on N ranks and load it on M ranks? Sure you can!
+    Activation checkpointing and torch-compile can also be different between save and load, no problem.
+    """
+    ckpt_dir = Path(ckpt_dir)
+    to_load = {"iteration": None}
+    to_load["model"] = dcpsd.get_model_state_dict(model)
+    if optimizer is not None:
+        to_load["optimizer"] = dcpsd.get_optimizer_state_dict(model, optimizer)
+    to_load.update(others)
+    dcp.load(
+        to_load,
+        storage_reader=dcpfs.FileSystemReader(ckpt_dir),
+        planner=dcp.default_planner.DefaultLoadPlanner(allow_partial_load=not strict_loading),
+        process_group=process_group,
+    )
+    iteration = to_load["iteration"]
+    dcpsd.set_model_state_dict(model, to_load["model"])
+    if optimizer is not None:
+        dcpsd.set_optimizer_state_dict(model, optimizer, to_load["optimizer"])
+    logger.info(f"Loaded: {ckpt_dir}")
+    return iteration
+
+
+def register_dont_save_hooks(module: torch.nn.Module, dont_save: Sequence[str]):
+    """
+    Registers save/load state dict hooks such that the weights in `dont_save` are not persisted in the checkpoint.
+
+    Typical use case: a classification model composed of a frozen backbone and a trainable head.
+    If the frozen backbone is loaded from torch hub, it does't make sense to save a copy of it in each checkpoint.
+    """
+
+    def state_dict_post_hook(module, state_dict, prefix, local_metadata):
+        # Remove frozen weights so they won't get saved.
+        # If this module is not the top-level module, its weights will have a prefix in the state dict.
+        nonlocal _dont_save
+        for k in _dont_save:
+            del state_dict[prefix + k]
+
+    def load_state_dict_pre_hook(
+        module,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        # This pre hook exists only to pass the prefix to the post hook when loading the state dict.
+        nonlocal _prefix
+        assert _prefix is None
+        _prefix = prefix
+
+    def load_state_dict_post_hook(module, incompatible_keys):
+        # Remove the frozen weights from the missing keys so they don't raise an error.
+        nonlocal _prefix
+        assert _prefix is not None
+        to_remove = []
+        for missing_key in incompatible_keys.missing_keys:
+            k = missing_key.removeprefix(_prefix)
+            k = k.replace("_checkpoint_wrapped_module.", "")  # Added by activation checkpointing
+            if k in _dont_save:
+                to_remove.append(missing_key)
+        for r in to_remove:
+            incompatible_keys.missing_keys.remove(r)
+        _prefix = None
+
+    _dont_save = set(name.replace("_checkpoint_wrapped_module.", "") for name in dont_save)
+    _prefix = None
+    module.register_state_dict_post_hook(state_dict_post_hook)
+    module.register_load_state_dict_pre_hook(load_state_dict_pre_hook)
+    module.register_load_state_dict_post_hook(load_state_dict_post_hook)
+
+
+def find_all_checkpoints(ckpt_dir: Path | str) -> list[Path]:
+    """Find all checkpoints in a directory, i.e. subdirs with integer name. Sorted from first to last."""
+    ckpt_dir = Path(ckpt_dir)
+    if not ckpt_dir.is_dir():
+        return []
+    checkpoints = [p for p in ckpt_dir.iterdir() if p.is_dir() and _is_int(p.name)]
+    checkpoints.sort(key=lambda p: int(p.name))
+    return checkpoints
+
+
+def find_latest_checkpoint(ckpt_dir: Path | str) -> Path | None:
+    """Find the latest checkpoint in a directory, i.e. the subdir with the highest integer name."""
+    checkpoints = find_all_checkpoints(ckpt_dir)
+    if len(checkpoints) == 0:
+        return None
+    return checkpoints[-1]
+
+
+def keep_last_n_checkpoints(ckpt_dir: Path | str, n: int | None):
+    """In a directory with integer-named subdirs, keep only the n subdirs with the highest number."""
+    if n is None:
+        return
+    checkpoints = find_all_checkpoints(ckpt_dir)
+    for ckpt_dir in checkpoints[:-n]:
+        try:
+            shutil.rmtree(ckpt_dir)
+            logger.info(f"Deleted: {ckpt_dir}")
+        except Exception:
+            logger.exception(f"Failed to delete: {ckpt_dir}")
+
+
+def keep_checkpoint_copy(src: Path | str):
+    """Copy a file/directory next to itself with a _keep suffix. Files are hardlinked."""
+    src = Path(src)
+    dst = src.parent / f"{src.name}_keep"
+    subprocess.check_output(["cp", "--recursive", "--link", src, dst])
+    logger.info(f"Copied: {src} -> {dst}")
+
+
+def _is_int(s: str) -> bool:
+    try:
+        int(s)
+        return True
+    except ValueError:
+        return False
+
+
+# Initialize a FSDP2 model from DCP or PyTorch standard checkpoint
+def init_fsdp_model_from_checkpoint(
+    model: torch.nn.Module,
+    checkpoint_path: str,
+    skip_load_prefixes: List[str] | None = None,
+    prefixes_not_sharded: List[str] | None = None,
+    process_group: dist.ProcessGroup = None,
+):
+    if not Path(checkpoint_path).is_dir():  # PyTorch standard checkpoint
+        logger.info(f"Loading pretrained weights from {checkpoint_path}")
+        chkpt = torch.load(checkpoint_path, map_location="cpu")["teacher"]
+        from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+
+        if process_group is None:
+            world_mesh = init_device_mesh(
+                "cuda",
+                mesh_shape=(dist.get_world_size(),),
+                mesh_dim_names=("dp",),
+            )
+        else:
+            world_mesh = DeviceMesh.from_group(process_group, "cuda")
+        chkpt = {
+            k: (
+                torch.distributed.tensor.distribute_tensor(v, world_mesh, src_data_rank=None)
+                if not k.startswith(pns)
+                else v
+            )
+            for pns in prefixes_not_sharded
+            for k, v in chkpt.items()
+        }
+        model.load_state_dict(
+            {k: v for k, v in chkpt.items() if not any(k.startswith(prefix) for prefix in skip_load_prefixes)}
+        )
+    else:  # DCP checkpoint
+        load_checkpoint(ckpt_dir=checkpoint_path, model=model, process_group=process_group)
+
+
+# Initialize a standard non distributed PyTorch model from PyTorch standard checkpoint for evals
+def init_model_from_checkpoint_for_evals(
+    model: torch.nn.Module, pretrained_weights: str | Path, checkpoint_key: str = None
+):
+    state_dict = torch.load(pretrained_weights, map_location="cpu")
+    if checkpoint_key is not None and checkpoint_key in state_dict:
+        logger.info(f"Take key {checkpoint_key} in provided checkpoint dict")
+        state_dict = state_dict[checkpoint_key]
+    # remove `module.` prefix
+    state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+    # remove `backbone.` prefix induced by multicrop wrapper
+    state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+    msg = model.load_state_dict(state_dict, strict=False)
+    logger.info("Pretrained weights found at {} and loaded with msg: {}".format(pretrained_weights, msg))
+
+
+def cleanup_checkpoint(ckpt_dir: str, checkpoint_retention_policy: CheckpointRetentionPolicy):
+    """
+    ckpt_dir is the directory containing each individual checkpoint directories (either at iteration, best (validation performance) or final)
+    |-- ckpt_dir/
+    |   |-- 0/
+    |       |--checkpoint.pth  or dcp_sharded_checkpoint_dir
+    |   |-- 99/
+            |--checkpoint.pth or dcp_sharded_checkpoint_dir
+    |   |-- 199/
+            |--checkpoint.pth or dcp_sharded_checkpoint_dir
+    |   |-- best/
+            |--checkpoint.pth or dcp_sharded_checkpoint_dir
+    |   |-- 299/
+            |--checkpoint.pth or dcp_sharded_checkpoint_dir
+    |   |-- final/
+            |--checkpoint.pth or dcp_sharded_checkpoint_dir
+    """
+    ckpt_dir = Path(ckpt_dir)
+    if not ckpt_dir.is_dir():
+        return []
+    checkpoint_filters = checkpoint_retention_policy.keep_filters
+    checkpoints = [p for p in ckpt_dir.iterdir() if p.is_dir()]
+    for checkpoint in checkpoints:
+        if checkpoint in checkpoint_filters:
+            continue
+        try:
+            shutil.rmtree(checkpoint)
+            logger.info(f"Deleted: {checkpoint}")
+        except Exception:
+            logger.exception(f"Failed to delete: {checkpoint}")

InfiniDepth/model/block/torchhub/dinov3/dinov3/configs/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from .config import (
+    DinoV3SetupArgs,
+    apply_scaling_rules_to_cfg,
+    exit_job,
+    get_cfg_from_args,
+    get_default_config,
+    setup_config,
+    setup_job,
+    setup_multidistillation,
+    write_config,
+)

InfiniDepth/model/block/torchhub/dinov3/dinov3/configs/config.py

@@ -0,0 +1,222 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+import math
+import os
+import pathlib
+import sys
+from dataclasses import dataclass, field
+from datetime import timedelta
+from typing import Any, List, Optional, Sequence, Tuple
+
+from omegaconf import DictConfig, OmegaConf
+
+import dinov3.distributed as distributed
+from dinov3.logging import cleanup_logging, setup_logging
+from dinov3.utils import fix_random_seeds, get_conda_env, get_sha
+
+logger = logging.getLogger("dinov3")
+
+
+@dataclass
+class DinoV3SetupArgs:
+    config_file: str
+    pretrained_weights: str | None = None
+    shard_unsharded_model: bool = False
+    output_dir: str = ""
+    opts: List[Any] = field(default_factory=lambda: [])
+
+    def __post_init__(self):
+        # When loaded from benchmark.yaml, self.opts is a frozen omegaconf.ListConfig,
+        # which works everywhere except when we want to modify it or when
+        # we try to json-serialize it. So we convert it to a regular list here.
+        if OmegaConf.is_config(self.opts):
+            self.opts = OmegaConf.to_object(self.opts)
+
+
+def apply_scaling_rules_to_cfg(cfg):  # to fix
+    assert distributed.is_enabled(), "Setup distributed to get global size !"
+    if "schedules" in cfg:
+        # For schedules v2, the scaling rules are applied when building the schedules, the config is not modified
+        return cfg
+
+    if cfg.optim.scaling_rule == "linear_wrt_256":
+        old_lr = cfg.optim.lr
+        cfg.optim.lr *= cfg.train.batch_size_per_gpu * distributed.get_world_size() / 256.0
+        logger.info(f"linear scaling learning rate; old: {old_lr}, new: {cfg.optim.lr}")
+    elif cfg.optim.scaling_rule == "sqrt_wrt_1024":
+        old_lr = cfg.optim.lr
+        cfg.optim.lr *= 4 * math.sqrt(cfg.train.batch_size_per_gpu * distributed.get_world_size() / 1024.0)
+        logger.info(f"sqrt scaling learning rate; old: {old_lr}, new: {cfg.optim.lr}")
+    return cfg
+
+
+def write_config(cfg, output_dir, name="config.yaml"):
+    logger.info(OmegaConf.to_yaml(cfg))
+    output_dir = os.path.abspath(output_dir)
+    saved_cfg_path = os.path.join(output_dir, name)
+    with open(saved_cfg_path, "w") as f:
+        OmegaConf.save(config=cfg, f=f)
+    return saved_cfg_path
+
+
+def get_default_config() -> DictConfig:
+    p = pathlib.Path(__file__).parent / "ssl_default_config.yaml"
+    return OmegaConf.load(p)
+
+
+def get_cfg_from_args(args: DinoV3SetupArgs, multidistillation=False, strict=True):
+    overrides = [*args.opts]
+    if args.output_dir is not None:
+        overrides.append(f"train.output_dir={os.path.realpath(args.output_dir)}")
+
+    # Config file
+    cfg = OmegaConf.load(args.config_file)
+
+    # Command line overrides
+    opts_cfg = OmegaConf.from_cli(overrides)
+
+    if multidistillation:
+        cfg = OmegaConf.merge(cfg, opts_cfg)
+    else:
+        # Default config
+        default_cfg = get_default_config()
+        if strict:
+            OmegaConf.set_struct(default_cfg, True)
+        cfg = OmegaConf.merge(default_cfg, cfg, opts_cfg)
+    return cfg
+
+
+def setup_config(args: DinoV3SetupArgs, strict_cfg=True):
+    """
+    Create configs and perform basic setups.
+    """
+    # Create the cfg with OmegaConf
+    cfg = get_cfg_from_args(args, strict=strict_cfg)
+    # setup distributed, logging, and random seeds
+    logger.info("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+    # dump config before modifying so it can be reloaded
+    if args.output_dir is not None:
+        write_config(cfg, args.output_dir)
+    # modify the config inplace by applying scaling rules
+    apply_scaling_rules_to_cfg(cfg)
+    return cfg
+
+
+def _enumerate_all_subgroup_ranks(all_subgroup_rank_spans: Sequence[Tuple[int, int]]):
+    """Expands a specification of process subgroups from spans to enumerated ranks.
+
+    Args:
+       all_group_rank_spans: a sequence of rank spans (first rank, last rank),
+           one for each process group. Example: ((0, 1), (2, 3), (4, 7)).
+    """
+    for first, last in all_subgroup_rank_spans:
+        assert first <= last
+    return tuple(tuple(range(first, last + 1)) for first, last in all_subgroup_rank_spans)
+
+
+def setup_multidistillation(args: DinoV3SetupArgs):
+    base_output_dir = args.output_dir
+    os.makedirs(args.output_dir, exist_ok=True)
+    # get config file for this rank
+    base_cfg = OmegaConf.load(args.config_file)
+    assert base_cfg.multidistillation.enabled
+
+    global_batch_size = base_cfg.multidistillation.global_batch_size
+
+    distributed.enable(overwrite=True)
+    seed = getattr(args, "seed", 0)
+    rank = distributed.get_rank()
+
+    # build process subgroups
+    all_subgroup_rank_spans = tuple(
+        (student.ranks_range[0], student.ranks_range[1] - 1) for student in base_cfg.multidistillation.students
+    )
+    all_subgroup_ranks = _enumerate_all_subgroup_ranks(all_subgroup_rank_spans)
+    distributed.new_subgroups(all_subgroup_ranks)
+
+    found = False
+    for student in base_cfg.multidistillation.students:
+        if rank in range(*student.ranks_range):
+            found = True
+            break
+    assert found, "rank of worker not in defined range"
+
+    name = student.name
+    config_path = student.config_path
+    n_gpus = student.ranks_range[1] - student.ranks_range[0]
+    assert global_batch_size % n_gpus == 0
+    total_n_gpus = distributed.get_world_size()
+
+    args.output_dir = os.path.join(base_output_dir, name)
+    args.opts += [f"train.output_dir={args.output_dir}"]
+    args.opts += [f"train.batch_size_per_gpu={global_batch_size // total_n_gpus}"]
+    args.config_file = os.path.abspath(config_path)
+    default_cfg = get_default_config()
+    cfg = OmegaConf.load(args.config_file)
+    cfg = OmegaConf.merge(default_cfg, cfg, base_cfg, OmegaConf.from_cli(args.opts))
+
+    global logger
+    setup_logging(output=args.output_dir, level=logging.INFO)
+
+    fix_random_seeds(seed + rank)
+
+    write_config(cfg, args.output_dir)
+    apply_scaling_rules_to_cfg(cfg)
+
+    return cfg
+
+
+def setup_job(
+    output_dir: Optional[str] = None,
+    distributed_enabled: bool = True,
+    logging_enabled: bool = True,
+    seed: Optional[int] = 0,
+    restrict_print_to_main_process: bool = True,
+    distributed_timeout: timedelta | None = None,
+):
+    """
+    Setup methods that should be done in every fairvit job
+    Initializes logging, distributed, random seeds and other utilities.
+    """
+    if output_dir is not None:
+        output_dir = os.path.realpath(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+
+    if logging_enabled:
+        setup_logging(
+            output=output_dir,
+            level=logging.INFO,
+            log_to_stdout_only_in_main_process=restrict_print_to_main_process,
+        )
+
+    if distributed_enabled:
+        distributed.enable(
+            overwrite=True,
+            nccl_async_error_handling=True,
+            restrict_print_to_main_process=restrict_print_to_main_process,
+            timeout=distributed_timeout,
+        )
+
+    if seed is not None:
+        rank = distributed.get_rank()
+        fix_random_seeds(seed + rank)
+
+    logger = logging.getLogger("fairvit")
+    logger.info("git:\n  {}\n".format(get_sha()))
+
+    # Log some python info
+    conda_env_name, conda_env_path = get_conda_env()
+    logger.info(f"conda env name: {conda_env_name}")
+    logger.info(f"conda env path: {conda_env_path}")
+    logger.info(f"python path: {sys.path}")
+
+
+def exit_job(distributed_enabled: bool = True, logging_enabled: bool = True):
+    if distributed_enabled:
+        distributed.disable()
+    if logging_enabled:
+        cleanup_logging()

InfiniDepth/model/block/torchhub/dinov3/dinov3/configs/ssl_default_config.yaml

@@ -0,0 +1,205 @@
+MODEL:
+  META_ARCHITECTURE: SSLMetaArch
+  DEVICE: cuda
+  WEIGHTS: ''
+  DTYPE: float32
+compute_precision:
+  param_dtype: bf16
+  reduce_dtype: fp32
+  sharding_strategy: SHARD_GRAD_OP
+dino:
+  loss_weight: 1.0
+  global_ignore_diagonal: true  # Whether to ignore A-A and B-B global pairs, default as in DINOv2, ignored by SSLMetaArch
+  head_n_prototypes: 65536
+  head_bottleneck_dim: 256
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 2048
+  koleo_loss_weight: 0.1
+  koleo_loss_distributed: false
+  koleo_topk: 1
+  koleo_distributed_replicas: 0
+  koleo_distributed_loss_group_size: null  # Size of the nearest neighbor set for distributed Koleo. If None, uses global batch size.
+  koleo_distributed_loss_group_data: true # group data from adjacent ranks to make sure koleo is applied on the same data distribution
+  force_weight_norm: false
+  reweight_dino_local_loss: false # If true, reweighting of DINO loss
+  local_loss_weight_schedule:  # Schedule for local loss weight, enabled if reweight_dino_local_loss is true
+    start: 0.5
+    peak: 0.5
+    end: 0.5
+    warmup_epochs: 0
+ibot:
+  loss_weight: 1.0
+  mask_sample_probability: 0.5
+  mask_ratio_min_max:
+  - 0.1
+  - 0.5
+  mask_random_circular_shift: false
+  force_masking_even_with_zero_weight: False
+  separate_head: true
+  head_n_prototypes: 65536
+  head_bottleneck_dim: 256
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 2048
+gram:
+  use_loss: false # (bool) if true gram is used, else not
+  compute_stats: false # (bool): if true compute auxilliary stats
+  loss_weight: 1.0 # (float): weight of the loss
+  ema_teacher: false # (bool): using the EMA teacher as GRAM teacher
+  ckpt: null  #(str): Checkpoint to the teacher
+  it_load_ema_teacher: -1 # (int): iteration at which the ema teacher is loaded into the gram teacher
+  rep_update: true # (bool): if true GRAM teacher updated every gram.update_frequency after iter gram.it_first_update steps
+  update_frequency: 50000 # (int): update frequency
+  it_first_update: 0 # (int): iteration of the first update
+  max_updates: null # (int): maximum number of updates to gram teacher. If None, it is unlimited
+  normalized: true # (bool): normalization of the features
+  img_level: false # (bool): if true GRAM computation at the image else, otherwise at the local batch level
+  remove_neg: false # (bool): if true remove the negative similarities before applying the loss
+  remove_only_teacher_neg: false # (bool): remove negative similarities of the teacher
+  tokens_used: all # (str): In [all, masked, unmasked]
+  global_teacher_resize_method: bicubic  # Method for resizing the outputs of the gram teacher
+  global_teacher_resize_antialias: false  # Whether to use antialiasing when resizing the outputs of the gram teacher
+  loss_weight_schedule: null  # (dict): If not None, use a schedule for the loss weight instead of `loss_weight`
+train:
+  batch_size_per_gpu: 64
+  dataset_path: ImageNet:split=TRAIN
+  data_config: null
+  output_dir: .
+  saveckp_freq: 20
+  seed: 0
+  num_workers: 10
+  OFFICIAL_EPOCH_LENGTH: 1250
+  monitor_gradient_norm: false
+  chunk_schedule: []
+  use_teacher_head: true
+  learn_from_teacher_tokens: false
+  centering: "sinkhorn_knopp" # or "sinkhorn_knopp"
+  checkpointing: false
+  checkpointing_full: false  # aggressive checkpointing
+  compile: true
+  cudagraphs: false
+  sharded_eval_checkpoint: false
+  cache_dataset: false
+student:
+  arch: vit_large
+  patch_size: 16
+  drop_path_rate: 0.3
+  layerscale: 1.0e-05
+  pretrained_weights: ''
+  ffn_layer: "mlp"
+  ffn_ratio: 4.0
+  resume_from_teacher_chkpt: ""
+  qkv_bias: true
+  proj_bias: true
+  ffn_bias: true
+  norm_layer: "layernorm"
+  n_storage_tokens: 0
+  mask_k_bias: false
+  untie_cls_and_patch_norms: false  # If true, use separate norms for CLS/reg and patch/mask tokens
+  untie_global_and_local_cls_norm: false  # If true, use separate norms for local and global crop CLS token during training
+  in_chans: 3
+  pos_embed_type: rope
+  pos_embed_rope_base: 100.0
+  pos_embed_rope_min_period: null
+  pos_embed_rope_max_period: null
+  pos_embed_rope_normalize_coords: separate  # min, max, separate
+  pos_embed_rope_shift_coords: null
+  pos_embed_rope_jitter_coords: null
+  pos_embed_rope_rescale_coords: null
+  pos_embed_rope_dtype: bf16
+  fp8_enabled: False  # Convert Linear layers to operate in fp8 precision
+  fp8_filter: "blocks"  # Regex that must appear in module path; empty means everything
+teacher:
+  momentum_teacher: 0.992
+  final_momentum_teacher: 1
+  warmup_teacher_temp: 0.04
+  teacher_temp: 0.07
+  warmup_teacher_temp_epochs: 30
+  in_chans: 3
+distillation:  # teacher
+  enabled: false
+  full_cfg_path: ""
+  checkpoint_path: ""
+multidistillation:
+  enabled: false
+hrft:  # non-hrft'd student
+  enabled: false
+  checkpoint_path: ""  # teacher_checkpoint path
+optim:
+  epochs: 100
+  optimizer: adamw
+  weight_decay: 0.04
+  weight_decay_end: 0.4
+  lr: 0.001
+  warmup_epochs: 10
+  min_lr: 1.0e-06
+  schedule_trunc_extra: 0.0  # Compute the schedule for (1 + schedule_trunc_extra) steps and truncate, .25 is a good choice
+  clip_grad: 3.0
+  freeze_last_layer_epochs: 1
+  scaling_rule: sqrt_wrt_1024
+  patch_embed_lr_mult: 0.2
+  dino_head_wd_multiplier: 1.0
+  layerwise_decay: 0.9
+  multi_tensor_optim: true
+  dump_fsdp_weights_path: ""
+  adamw_beta1: 0.9
+  adamw_beta2: 0.999
+crops:
+  global_crops_scale:
+  - 0.32
+  - 1.0
+  local_crops_number: 8
+  local_crops_scale:
+  - 0.05
+  - 0.32
+  global_crops_size: 224
+  local_crops_size: 96
+  global_local_crop_pairs_ratios: 1.0
+  gram_teacher_crops_size: null  # If not None, return crops for gram teacher
+  localcrops_subset_of_globalcrops: false
+  share_color_jitter: false
+  horizontal_flips: true
+  gram_teacher_no_distortions: false  # If True, no distortions are applied to gram teacher crops
+  rgb_mean:
+  - 0.485
+  - 0.456
+  - 0.406
+  rgb_std:
+  - 0.229
+  - 0.224
+  - 0.225
+evaluation:
+  eval_period_iterations: 12500
+  low_freq_every: 5
+  config_files:  # Must be in fairvit/eval/configs
+    high_freq: benchmark_high_frequency.yaml  # More often
+    low_freq: benchmark_low_frequency.yaml  # Less often
+checkpointing:
+  period: 3750
+  max_to_keep: 3
+  keep_every: 99999999999999999  # Save a checkpoint every N iterations, regardless of max_to_keep and period
+
+# Example of constant schedules with schedules v2
+# # schedules:
+# #   lr:
+# #     start: 0.0
+# #     peak: 1e-3
+# #     end: 1e-6
+# #     warmup_epochs: 10
+# #     freeze_last_layer_epochs: 1
+# #   weight_decay:
+# #     start: 0.04
+# #     peak: 0.04
+# #     end: 0.04
+# #     warmup_epochs: 0
+# #   momentum:
+# #     start: 0.992
+# #     peak: 0.992
+# #     end: 0.992
+# #     warmup_epochs: 0
+# #   teacher_temp:
+# #     start: 0.04
+# #     peak: 0.07
+# #     end: 0.07
+# #     warmup_epochs: 30