init!

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+cache/
+output/

app.py

@@ -0,0 +1,165 @@
+import os
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import (
+    InitProcessGroupKwargs, 
+    ProjectConfiguration, 
+    set_seed
+)
+import torch
+from contextlib import nullcontext
+import trimesh
+import gradio as gr
+from gradio_imageslider import ImageSlider
+from da2.utils.base import load_config
+from da2.utils.model import load_model
+from da2.utils.io import (
+    read_cv2_image,
+    torch_transform,
+    tensorize
+)
+from da2.utils.vis import colorize_distance
+from da2.utils.d2pc import distance2pointcloud
+from datetime import (
+    timedelta,
+    datetime
+)
+import cv2
+import numpy as np
+
+last_glb_path = None
+
+def prepare_to_run_demo():
+    config = load_config('configs/infer.json')
+    kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=config['accelerator']['timeout']))
+    output_dir = f'output/infer'
+    if not os.path.exists(output_dir): os.makedirs(output_dir, exist_ok=True)
+    accu_steps = config['accelerator']['accumulation_nsteps']
+    accelerator = Accelerator(
+        gradient_accumulation_steps=accu_steps,
+        mixed_precision=config['accelerator']['mixed_precision'],
+        log_with=config['accelerator']['report_to'],
+        project_config=ProjectConfiguration(project_dir=output_dir),
+        kwargs_handlers=[kwargs]
+    )
+    logger = get_logger(__name__, log_level='INFO')
+    config['env']['logger'] = logger
+    set_seed(config['env']['seed'])
+    return config, accelerator
+
+def read_mask_demo(mask_path, shape):
+    if mask_path is None:
+        return np.ones(shape[1:]) > 0
+    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
+    mask = mask > 0
+    return mask
+
+def load_infer_data_demo(image, mask, model_dtype, device):
+    cv2_image = read_cv2_image(image)
+    image = torch_transform(cv2_image)
+    mask = read_mask_demo(mask, image.shape)
+    image = tensorize(image, model_dtype, device)
+    return image, cv2_image, mask
+
+def ply2glb(ply_path, glb_path):
+    pcd = trimesh.load(ply_path)
+    points = np.asarray(pcd.vertices)
+    colors = np.asarray(pcd.visual.vertex_colors)
+    cloud = trimesh.points.PointCloud(vertices=points, colors=colors)
+    cloud.export(glb_path)
+    os.remove(ply_path)
+
+def fn(image_path, mask_path):
+    global last_glb_path
+    config, accelerator = prepare_to_run_demo()
+    model = load_model(config, accelerator)
+    image, cv2_image, mask = load_infer_data_demo(image_path, mask_path, 
+        model_dtype=config['spherevit']['dtype'], device=accelerator.device)
+    if torch.backends.mps.is_available():
+        autocast_ctx = nullcontext()
+    else:
+        autocast_ctx = torch.autocast(accelerator.device.type)
+    with autocast_ctx, torch.no_grad():
+        distance = model(image).cpu().numpy()[0]
+        if last_glb_path is not None:
+            os.remove(last_glb_path)
+        distance_vis = colorize_distance(distance, mask)
+        save_path = f'cache/tmp_{datetime.now().strftime("%Y%m%d_%H%M%S")}.glb'
+        last_glb_path = save_path
+        normal_image = distance2pointcloud(distance, cv2_image, mask, save_path=save_path.replace('.glb', '.ply'), return_normal=True, save_distance=False)
+        ply2glb(save_path.replace('.glb', '.ply'), save_path)
+        return save_path, [distance_vis, normal_image]
+
+inputs = [
+    gr.Image(label="Input Image", type="filepath"),
+    gr.Image(label="Input Mask", type="filepath"),
+]
+outputs = [
+    gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Point Cloud"),
+    gr.ImageSlider(
+        label="Output Depth / Normal (transformed from the depth)",
+        type="pil",
+        slider_position=75,
+    )
+]
+
+demo = gr.Interface(
+    fn=fn,
+    title="DA<sup>2</sup>: <u>D</u>epth <u>A</u>nything in <u>A</u>ny <u>D</u>irection",
+    description="""
+        <p align="center">
+        <a title="Project Page" href="https://depth-any-in-any-dir.github.io/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <img src="https://img.shields.io/badge/Project-Website-pink?logo=googlechrome&logoColor=white">
+        </a>
+        <a title="arXiv" href="http://arxiv.org/abs/2509.26618" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <img src="https://img.shields.io/badge/arXiv-Paper-b31b1b?logo=arxiv&logoColor=white">
+        </a>
+        <a title="Github" href="https://github.com/EnVision-Research/DA-2" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <img src="https://img.shields.io/github/stars/EnVision-Research/DA-2?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
+        </a>
+        <a title="Social" href="https://x.com/_akhaliq/status/1973283687652606411" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
+        </a>
+        <a title="Social" href="https://x.com/haodongli00/status/1973287870317338747" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
+        </a>
+        <br>
+        <strong>Please consider starring <span style="color: orange">&#9733;</span> our <a href="https://github.com/EnVision-Research/DA-2" target="_blank" rel="noopener noreferrer">GitHub Repo</a> if you find this demo useful!</strong>
+        </p>
+        <p><strong>Note: the "Input Mask" is optional, all pixels are assumed to be valid if mask is None.</strong></p>
+    """,
+    inputs=inputs,
+    outputs=outputs,
+    examples=[
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a1.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a2.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a3.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a4.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/b0.png"), 
+         os.path.join(os.path.dirname(__file__), "assets/masks/b0.png")],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/b1.png"), 
+         os.path.join(os.path.dirname(__file__), "assets/masks/b1.png")],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a5.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a6.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a7.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a8.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/b2.png"), 
+         os.path.join(os.path.dirname(__file__), "assets/masks/b2.png")],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/b3.png"), 
+         os.path.join(os.path.dirname(__file__), "assets/masks/b3.png")],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a9.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a10.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a11.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/a0.png"), None],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/b4.png"), 
+         os.path.join(os.path.dirname(__file__), "assets/masks/b4.png")],
+        [os.path.join(os.path.dirname(__file__), "assets/demos/b5.png"), 
+         os.path.join(os.path.dirname(__file__), "assets/masks/b5.png")],
+    ],
+    examples_per_page=20
+)
+
+demo.launch(
+        server_name="0.0.0.0",
+        server_port=6381,
+)

configs/accelerate/0.yaml

@@ -0,0 +1,16 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: 'NO'
+downcast_bf16: 'no'
+gpu_ids: '0'
+machine_rank: 0
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

configs/infer.json

@@ -0,0 +1,39 @@
+{
+    "env": {
+        "seed": 42,
+        "verbose": true
+    },
+    "accelerator": {
+        "report_to": ["tensorboard"],
+        "mixed_precision": "fp16",
+        "accumulation_nsteps": 4,
+        "timeout": 36000
+    },
+    "inference": {
+        "images": "assets/demos",
+        "masks": "assets/masks",
+        "min_pixels": 580000,
+        "max_pixels": 620000
+    },
+    "spherevit": {
+        "vit_w_esphere": {
+            "input_dims": [1024, 1024, 1024, 1024],
+            "hidden_dim": 512,
+            "num_heads": 8,
+            "expansion": 4,
+            "num_layers_head": [2, 2, 2],
+            "dropout": 0.0,
+            "layer_scale": 0.0001,
+            "out_dim": 64,
+            "kernel_size": 3,
+            "num_prompt_blocks": 1,
+            "use_norm": false
+        },
+        "sphere": {
+            "width": 1092,
+            "height": 546,
+            "hfov": 6.2832,
+            "vfov": 3.1416
+        }
+    }
+}

requirements.txt

@@ -0,0 +1 @@
+pip install -e src/

src/da2.egg-info/PKG-INFO

@@ -0,0 +1,23 @@
+Metadata-Version: 2.4
+Name: da2
+Version: 0.1.0
+Summary: For the implementation of DA^2: Depth Anything in Any Direction
+Author-email: "H. Li" <hal211@ucsd.edu>
+Requires-Dist: torch==2.5.0
+Requires-Dist: torchvision==0.20.0
+Requires-Dist: torchaudio==2.5.0
+Requires-Dist: xformers==0.0.28.post2
+Requires-Dist: diffusers==0.32.0
+Requires-Dist: tensorboard==2.18.0
+Requires-Dist: utils3d@ git+https://github.com/EasternJournalist/utils3d.git@3913c65d81e05e47b9f367250cf8c0f7462a0900
+Requires-Dist: opencv-python==4.12.0.88
+Requires-Dist: gradio==5.49.0
+Requires-Dist: gradio-client==1.13.3
+Requires-Dist: gradio-imageslider==0.0.20
+Requires-Dist: accelerate==1.1.1
+Requires-Dist: omegaconf==2.3.0
+Requires-Dist: tabulate==0.9.0
+Requires-Dist: einops==0.8.0
+Requires-Dist: timm==1.0.15
+Requires-Dist: trimesh==4.5.2
+Requires-Dist: transformers==4.46.3

src/da2.egg-info/SOURCES.txt

@@ -0,0 +1,28 @@
+pyproject.toml
+da2/__init__.py
+da2.egg-info/PKG-INFO
+da2.egg-info/SOURCES.txt
+da2.egg-info/dependency_links.txt
+da2.egg-info/requires.txt
+da2.egg-info/top_level.txt
+da2/model/__init__.py
+da2/model/base.py
+da2/model/sphere.py
+da2/model/spherevit.py
+da2/model/vit_w_esphere.py
+da2/model/dinov2/__init__.py
+da2/model/dinov2/attention.py
+da2/model/dinov2/block.py
+da2/model/dinov2/dino_head.py
+da2/model/dinov2/dinovit.py
+da2/model/dinov2/drop_path.py
+da2/model/dinov2/layer_scale.py
+da2/model/dinov2/mlp.py
+da2/model/dinov2/patch_embed.py
+da2/model/dinov2/swiglu_ffn.py
+da2/utils/__init__.py
+da2/utils/base.py
+da2/utils/d2pc.py
+da2/utils/io.py
+da2/utils/model.py
+da2/utils/vis.py

src/da2.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

src/da2.egg-info/requires.txt

@@ -0,0 +1,18 @@
+torch==2.5.0
+torchvision==0.20.0
+torchaudio==2.5.0
+xformers==0.0.28.post2
+diffusers==0.32.0
+tensorboard==2.18.0
+utils3d@ git+https://github.com/EasternJournalist/utils3d.git@3913c65d81e05e47b9f367250cf8c0f7462a0900
+opencv-python==4.12.0.88
+gradio==5.49.0
+gradio-client==1.13.3
+gradio-imageslider==0.0.20
+accelerate==1.1.1
+omegaconf==2.3.0
+tabulate==0.9.0
+einops==0.8.0
+timm==1.0.15
+trimesh==4.5.2
+transformers==4.46.3

src/da2.egg-info/top_level.txt

@@ -0,0 +1 @@
+da2

src/da2/__init__.py

@@ -0,0 +1,25 @@
+from .utils.base import (
+    prepare_to_run
+)
+from .utils.model import (
+    load_model
+)
+from .utils.io import (
+    load_infer_data
+)
+from .utils.vis import (
+    colorize_distance,
+    concatenate_images
+)
+from .utils.d2pc import (
+    distance2pointcloud
+)
+
+__all__ = [
+    'prepare_to_run',
+    'load_model',
+    'load_infer_data',
+    'colorize_distance',
+    'concatenate_images',
+    'distance2pointcloud'
+]

src/da2/model/__init__.py

@@ -0,0 +1,11 @@
+from .spherevit import (
+    SphereViT
+)
+from .vit_w_esphere import (
+    ViT_w_Esphere
+)
+
+__all__ = [
+    'SphereViT',
+    'ViT_w_Esphere',
+]

src/da2/model/base.py

@@ -0,0 +1,393 @@
+import torch
+import torch.nn as nn
+from math import log2, pi
+from typing import Tuple
+import torch.nn.functional as F
+from einops import rearrange
+from functools import partial
+
+
+def fourier_dimension_expansion(
+    x: torch.Tensor,
+    dim: int = 512,
+    max_freq: int = 64,
+    use_cos: bool = True,
+    use_log: bool = True,
+):
+    device, dtype, input_dim = x.device, x.dtype, x.shape[-1]
+    # input_dim: 2
+    num_bands = dim // (2 * input_dim) if use_cos else dim // input_dim
+    # num_bands = 512 // 2 = 256
+    if use_log:
+        scales = 2.0 ** torch.linspace(
+            0.0, log2(max_freq), steps=num_bands, device=device, dtype=dtype
+        )
+    else:
+        scales = torch.linspace(
+            1.0, max_freq / 2, num_bands, device=device, dtype=dtype
+        )
+    x = x.unsqueeze(-1)
+    scales = scales[(*((None,) * (len(x.shape) - 1)), Ellipsis)]
+    x = x * scales * pi
+    x = torch.cat(
+        (
+            [x.sin(), x.cos()]
+            if use_cos
+            else [
+                x.sin(),
+            ]
+        ),
+        dim=-1,
+    )
+    x = x.flatten(-2)
+    return x
+
+def flatten(
+    flat_tensor: torch.Tensor,
+    old: Tuple[int, int],
+    new: Tuple[int, int],
+) -> torch.Tensor:
+    if old[0] == new[0] and old[1] == new[1]:
+        return flat_tensor
+    tensor = flat_tensor.view(flat_tensor.shape[0], old[0], old[1], -1).permute(
+        0, 3, 1, 2
+    )  # b c h w
+    tensor_interp = F.interpolate(
+        tensor,
+        size=(new[0], new[1]),
+        mode='nearest',
+    )
+    flat_tensor_interp = tensor_interp.view(
+        flat_tensor.shape[0], -1, new[0] * new[1]
+    ).permute(
+        0, 2, 1
+    )  # b (h w) c
+    return flat_tensor_interp.contiguous()
+
+
+class DimensionAligner(nn.Module):
+    def __init__(self, input_dims: list[int], hidden_dim: int):
+        super().__init__()
+        self.aligners = nn.ModuleList([])
+        self.num_chunks = len(input_dims)
+        self.checkpoint = True
+        for input_dim in input_dims:
+            self.aligners.append(nn.Linear(input_dim, hidden_dim))
+
+    def forward(self, xs: torch.Tensor) -> torch.Tensor:
+        outs = [self.aligners[i](x) for i, x in enumerate(xs)]
+        return outs
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: float | torch.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: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if callable(d) else d
+
+
+class SwiGLU(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, gates = x.chunk(2, dim=-1)
+        return x * F.silu(gates)
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        expansion: int = 4,
+        dropout: float = 0.0,
+        gated: bool = False,
+        output_dim: int | None = None,
+    ):
+        super().__init__()
+        if gated:
+            expansion = int(expansion * 2 / 3)
+        hidden_dim = int(input_dim * expansion)
+        output_dim = default(output_dim, input_dim)
+        self.norm = nn.LayerNorm(input_dim)
+        self.proj1 = nn.Linear(input_dim, hidden_dim)
+        self.proj2 = nn.Linear(hidden_dim, output_dim)
+        self.act = nn.GELU() if not gated else SwiGLU()
+        self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.norm(x)
+        x = self.proj1(x)
+        x = self.act(x)
+        x = self.proj2(x)
+        x = self.dropout(x)
+        return x
+
+
+class AttentionBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 4,
+        expansion: int = 4,
+        dropout: float = 0.0,
+        cosine: bool = False,
+        gated: bool = False,
+        layer_scale: float = 1.0,
+        context_dim: int | None = None,
+        detach_query: bool = False,
+        residual_ls: bool = False,
+    ):
+        super().__init__()
+        self.dropout = dropout
+        self.num_heads = num_heads
+        self.hidden_dim = dim
+        context_dim = dim if context_dim is None else context_dim
+        self.mlp = MLP(dim, expansion=expansion, dropout=dropout, gated=gated)
+        self.kv = nn.Linear(context_dim, dim * 2, bias=False)
+        self.q = nn.Linear(dim, dim, bias=False)
+        self.norm_attnx = nn.LayerNorm(dim)
+        self.norm_attnctx = nn.LayerNorm(context_dim)
+        self.cosine = cosine
+        self.out = nn.Linear(dim, dim, bias=False)
+        self.ls1_1 = (
+            LayerScale(dim, layer_scale)
+            if layer_scale > 0.0 and not residual_ls
+            else nn.Identity()
+        )
+        self.ls1_2 = (
+            LayerScale(dim, layer_scale)
+            if layer_scale > 0.0 and residual_ls
+            else nn.Identity()
+        )
+        self.ls2 = LayerScale(dim, layer_scale) if layer_scale > 0.0 else nn.Identity()
+        self.detach_query = detach_query
+
+    def attn(
+        self,
+        x: torch.Tensor,
+        attn_bias: torch.Tensor | None = None,
+        context: torch.Tensor | None = None,
+        pos_embed: torch.Tensor | None = None,
+        pos_embed_context: torch.Tensor | None = None,
+        rope: nn.Module | None = None,
+        rope_pos: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.detach_query:
+            x = x.detach()
+        x = self.norm_attnx(x)
+        context = self.norm_attnctx(context)
+        k, v = rearrange(
+            self.kv(context), 'b n (kv h d) -> b h n d kv', h=self.num_heads, kv=2
+        ).unbind(dim=-1)
+        q = rearrange(self.q(x), 'b n (h d) -> b h n d', h=self.num_heads)
+
+        if rope is not None:
+            q = rope(q.permute(0, 2, 1, 3), input_pos=rope_pos).permute(0, 2, 1, 3)
+            k = rope(k.permute(0, 2, 1, 3), input_pos=rope_pos).permute(0, 2, 1, 3)
+        else:
+            if pos_embed is not None:
+                pos_embed = rearrange(
+                    pos_embed, 'b n (h d) -> b h n d', h=self.num_heads
+                )
+                q = q + pos_embed
+            if pos_embed_context is not None:
+                pos_embed_context = rearrange(
+                    pos_embed_context, 'b n (h d) -> b h n d', h=self.num_heads
+                )
+                k = k + pos_embed_context
+
+        if self.cosine:
+            q, k = map(partial(F.normalize, p=2, dim=-1), (q, k))  # cosine sim
+
+        x = F.scaled_dot_product_attention(
+            q, k, v, dropout_p=self.dropout, attn_mask=attn_bias
+        )
+        x = rearrange(x, 'b h n d -> b n (h d)')
+        x = self.out(x)
+        return x
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor | None = None,
+        pos_embed: torch.Tensor | None = None,
+        pos_embed_context: torch.Tensor | None = None,
+        attn_bias: torch.Tensor | None = None,
+        rope: nn.Module | None = None,
+        rope_pos: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        context = x if context is None else context
+        x = self.ls1_1(
+            self.attn(
+                x,
+                rope=rope,
+                rope_pos=rope_pos,
+                attn_bias=attn_bias,
+                context=context,
+                pos_embed=pos_embed,
+                pos_embed_context=pos_embed_context,
+            )
+        ) + self.ls1_2(x)
+        x = self.ls2(self.mlp(x)) + x
+        return x
+
+
+class AttentionSeq(nn.Module):
+    def __init__(
+        self,
+        num_blocks: int,
+        dim: int,
+        num_heads: int = 4,
+        expansion: int = 4,
+        dropout: float = 0.0,
+        cosine: bool = False,
+        gated: bool = False,
+        layer_scale: float = 1.0,
+        context_dim: int | None = None,
+        detach_query: bool = False,
+        residual_ls: bool = False,
+    ):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                AttentionBlock(
+                    dim=dim,
+                    num_heads=num_heads,
+                    expansion=expansion,
+                    dropout=dropout,
+                    cosine=cosine,
+                    gated=gated,
+                    layer_scale=layer_scale,
+                    context_dim=context_dim,
+                    detach_query=detach_query,
+                    residual_ls=residual_ls,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor | None = None,
+        pos_embed: torch.Tensor | None = None,
+        pos_embed_context: torch.Tensor | None = None,
+        attn_bias: torch.Tensor | None = None,
+        rope: nn.Module | None = None,
+        rope_pos: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        for layer in self.layers:
+            x = layer(
+                x,
+                context=context,
+                pos_embed=pos_embed,
+                pos_embed_context=pos_embed_context,
+                attn_bias=attn_bias,
+                rope=rope,
+                rope_pos=rope_pos,
+            )
+        return x
+
+
+class ResidualConvNet(nn.Module):
+    def __init__(
+        self,
+        dim,
+        kernel_size: int = 3,
+        padding_mode: str = 'zeros',
+        dilation: int = 1,
+        layer_scale: float = 1.0,
+        use_norm: bool = False,
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv2d(
+            dim,
+            dim,
+            kernel_size=kernel_size,
+            padding=dilation * (kernel_size - 1) // 2,
+            dilation=dilation,
+            padding_mode=padding_mode,
+        )
+        self.conv2 = nn.Conv2d(
+            dim,
+            dim,
+            kernel_size=kernel_size,
+            padding=dilation * (kernel_size - 1) // 2,
+            dilation=dilation,
+            padding_mode=padding_mode,
+        )
+        self.activation = nn.LeakyReLU()
+        self.gamma = (
+            nn.Parameter(layer_scale * torch.ones(1, dim, 1, 1))
+            if layer_scale > 0.0
+            else 1.0
+        )
+        self.norm1 = nn.GroupNorm(dim // 16, dim) if use_norm else nn.Identity()
+        self.norm2 = nn.GroupNorm(dim // 16, dim) if use_norm else nn.Identity()
+
+    def forward(self, x):
+        out = self.activation(x)
+        out = self.conv1(out)
+        out = self.norm1(out)
+        out = self.activation(out)
+        out = self.conv2(out)
+        out = self.norm2(out)
+        return self.gamma * out + x
+
+
+class ResidualUpsampler(nn.Module):
+    def __init__(
+        self,
+        hidden_dim,
+        output_dim: int = None,
+        num_layers: int = 2,
+        kernel_size: int = 3,
+        layer_scale: float = 1.0,
+        padding_mode: str = 'zeros',
+        use_norm: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        output_dim = output_dim if output_dim is not None else hidden_dim // 2
+        self.convs = nn.ModuleList([])
+        for _ in range(num_layers):
+            self.convs.append(
+                ResidualConvNet(
+                    hidden_dim,
+                    kernel_size=kernel_size,
+                    layer_scale=layer_scale,
+                    padding_mode=padding_mode,
+                    use_norm=use_norm,
+                )
+            )
+        self.up = nn.Sequential(
+            nn.Conv2d(
+                hidden_dim,
+                output_dim,
+                kernel_size=1,
+                padding=0,
+                padding_mode=padding_mode,
+            ),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+        )
+
+    def forward(self, x: torch.Tensor):
+        for conv in self.convs:
+            x = conv(x)
+        x = self.up(x)
+        return x

src/da2/model/dinov2/__init__.py

@@ -0,0 +1,13 @@
+# 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 .dinovit import (
+    DINOViT
+)
+
+__all__ = [
+    'DINOViT'
+]