# ruff: noqa # fmt: off """Depth Anything 3, vendored verbatim from ByteDance-Seed/depth-anything-3 @ 3d835ec1a5802d64a8b8b15f817a1ab54809bfe4: the files below concatenated, package-internal imports replaced by relative imports of the other vendored files (generated by scripts/model_ports/common/vendor_flat.py; do not edit by hand). Code license: Apache-2.0 (ByteDance).""" # ============================================================================================================== # depth_anything_3/utils/logger.py # ============================================================================================================== # Copyright (c) 2025 ByteDance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import sys class Color: RED = "\033[91m" YELLOW = "\033[93m" WHITE = "\033[97m" GREEN = "\033[92m" RESET = "\033[0m" LOG_LEVELS = {"ERROR": 0, "WARN": 1, "INFO": 2, "DEBUG": 3} COLOR_MAP = {"ERROR": Color.RED, "WARN": Color.YELLOW, "INFO": Color.WHITE, "DEBUG": Color.GREEN} def get_env_log_level(): level = os.environ.get("DA3_LOG_LEVEL", "INFO").upper() return LOG_LEVELS.get(level, LOG_LEVELS["INFO"]) class Logger: def __init__(self): self.level = get_env_log_level() def log(self, level_str, *args, **kwargs): level_key = level_str.split(":")[0].strip() level_val = LOG_LEVELS.get(level_key) if level_val is None: raise ValueError(f"Unknown log level: {level_str}") if self.level >= level_val: color = COLOR_MAP[level_key] msg = " ".join(str(arg) for arg in args) # Align log level output in square brackets # ERROR and DEBUG are 5 characters, INFO and WARN have an extra space for alignment tag = level_key if tag in ("INFO", "WARN"): tag += " " print( f"{color}[{tag}] {msg}{Color.RESET}", file=sys.stderr if level_key == "ERROR" else sys.stdout, **kwargs, ) def error(self, *args, **kwargs): self.log("ERROR:", *args, **kwargs) def warn(self, *args, **kwargs): self.log("WARN:", *args, **kwargs) def info(self, *args, **kwargs): self.log("INFO:", *args, **kwargs) def debug(self, *args, **kwargs): self.log("DEBUG:", *args, **kwargs) logger = Logger() __all__ = ["logger"] if __name__ == "__main__": logger.info("This is an info message") logger.warn("This is a warning message") logger.error("This is an error message") logger.debug("This is a debug message") # ============================================================================================================== # depth_anything_3/utils/constants.py # ============================================================================================================== # Copyright (c) 2025 ByteDance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. DEFAULT_MODEL = "depth-anything/DA3NESTED-GIANT-LARGE-1.1" DEFAULT_EXPORT_DIR = "workspace/gallery/scene" DEFAULT_GALLERY_DIR = "workspace/gallery" DEFAULT_GRADIO_DIR = "workspace/gradio" THRESH_FOR_REF_SELECTION = 3 # ============================================================================= # Benchmark Evaluation Constants # ============================================================================= # Default evaluation workspace directory DEFAULT_EVAL_WORKSPACE = "workspace/evaluation" # Default reference view selection strategy for evaluation # Use "first" for consistent and reproducible evaluation results # Other options: "saddle_balanced", "auto", "mid" EVAL_REF_VIEW_STRATEGY = "first" # ----------------------------------------------------------------------------- # DTU Dataset Configuration # Reference: https://roboimagedata.compute.dtu.dk/ # Note: DepthAnything3 was never trained on any images from DTU. # ----------------------------------------------------------------------------- # Root directory for DTU evaluation data (MVSNet format) # Download from: https://drive.google.com/file/d/1rX0EXlUL4prRxrRu2DgLJv2j7-tpUD4D/view DTU_EVAL_DATA_ROOT = "workspace/benchmark_dataset/dtu" # List of DTU evaluation scenes DTU_SCENES = [ "scan1", "scan4", "scan9", "scan10", "scan11", "scan12", "scan13", "scan15", "scan23", "scan24", "scan29", "scan32", "scan33", "scan34", "scan48", "scan49", "scan62", "scan75", "scan77", "scan110", "scan114", "scan118", ] # Point cloud fusion hyperparameters DTU_DIST_THRESH = 0.2 # Distance threshold for geometric consistency (mm) DTU_NUM_CONSIST = 4 # Minimum number of consistent views for a point DTU_MAX_POINTS = 4_000_000 # Maximum points in fused point cloud # 3D reconstruction evaluation hyperparameters DTU_DOWN_DENSE = 0.2 # Downsample density for evaluation (mm) DTU_PATCH_SIZE = 60 # Patch size for boundary handling DTU_MAX_DIST = 20 # Outlier threshold for accuracy/completeness (mm) # ----------------------------------------------------------------------------- # DTU-64 Dataset Configuration (Pose Evaluation Only) # This is a subset of DTU with 64 images per scene for pose evaluation. # Note: This dataset is ONLY for pose evaluation, not 3D reconstruction. # ----------------------------------------------------------------------------- # Root directory for DTU-64 evaluation data DTU64_EVAL_DATA_ROOT = "workspace/benchmark_dataset/dtu64" DTU64_CAMERA_ROOT = "workspace/benchmark_dataset/dtu64/Cameras" # List of DTU-64 evaluation scenes (13 scenes) DTU64_SCENES = [ "scan105", "scan114", "scan118", "scan122", "scan24", "scan37", "scan40", "scan55", "scan63", "scan65", "scan69", "scan83", "scan97", ] # ----------------------------------------------------------------------------- # ETH3D Dataset Configuration # Reference: https://www.eth3d.net/ # High-resolution multi-view stereo benchmark with laser-scanned ground truth. # Note: DepthAnything3 was never trained on any images from ETH3D. # ----------------------------------------------------------------------------- # Root directory for ETH3D evaluation data ETH3D_EVAL_DATA_ROOT = "workspace/benchmark_dataset/eth3d" # List of ETH3D evaluation scenes (indoor and outdoor) ETH3D_SCENES = [ "courtyard", "electro", "kicker", "pipes", "relief", # "terrace", # Excluded: known issues "delivery_area", "facade", # "meadow", # Excluded: known issues "office", "playground", "relief_2", "terrains", ] # Images to filter out (known problematic views per scene) ETH3D_FILTER_KEYS = { "delivery_area": ["711.JPG", "712.JPG", "713.JPG", "714.JPG"], "electro": ["9289.JPG", "9290.JPG", "9291.JPG", "9292.JPG", "9293.JPG", "9298.JPG"], "playground": ["587.JPG", "588.JPG", "589.JPG", "590.JPG", "591.JPG", "592.JPG"], "relief": [ "427.JPG", "428.JPG", "429.JPG", "430.JPG", "431.JPG", "432.JPG", "433.JPG", "434.JPG", "435.JPG", "436.JPG", "437.JPG", "438.JPG", ], "relief_2": [ "458.JPG", "459.JPG", "460.JPG", "461.JPG", "462.JPG", "463.JPG", "464.JPG", "465.JPG", "466.JPG", "467.JPG", "468.JPG", ], } # TSDF fusion hyperparameters (scaled for outdoor scenes) ETH3D_VOXEL_LENGTH = 4.0 / 512.0 * 5 # Voxel size for TSDF (meters) ETH3D_SDF_TRUNC = 0.04 * 5 # SDF truncation distance (meters) ETH3D_MAX_DEPTH = 100000.0 # Maximum depth for integration (effectively no truncation) # Point cloud sampling ETH3D_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh # 3D reconstruction evaluation hyperparameters ETH3D_EVAL_THRESHOLD = 0.05 * 5 # Distance threshold for precision/recall (meters) ETH3D_DOWN_SAMPLE = 4.0 / 512.0 * 5 # Voxel size for evaluation downsampling (meters) # ============================================================================== # 7Scenes Dataset Configuration # ============================================================================== # Reference: https://www.microsoft.com/en-us/research/project/rgb-d-dataset-7-scenes/ # Note: Indoor RGB-D dataset with ground truth poses and meshes. # Root directory for 7Scenes evaluation data SEVENSCENES_EVAL_DATA_ROOT = "workspace/benchmark_dataset/7scenes" # List of 7Scenes evaluation scenes SEVENSCENES_SCENES = [ "chess", "fire", "heads", "office", "pumpkin", "redkitchen", "stairs", ] # Fixed camera intrinsics for 7Scenes (all images share same intrinsics) SEVENSCENES_FX = 585.0 SEVENSCENES_FY = 585.0 SEVENSCENES_CX = 320.0 SEVENSCENES_CY = 240.0 # TSDF fusion hyperparameters (indoor scenes, smaller voxels) SEVENSCENES_VOXEL_LENGTH = 4.0 / 512.0 # Voxel size for TSDF (meters) SEVENSCENES_SDF_TRUNC = 0.04 # SDF truncation distance (meters) SEVENSCENES_MAX_DEPTH = 1000000.0 # Maximum depth for integration (no truncation) # Point cloud sampling SEVENSCENES_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh # 3D reconstruction evaluation hyperparameters SEVENSCENES_EVAL_THRESHOLD = 0.05 # Distance threshold for precision/recall (meters) SEVENSCENES_DOWN_SAMPLE = 4.0 / 512.0 # Voxel size for evaluation downsampling (meters) # ============================================================================== # ScanNet++ Dataset Configuration # ============================================================================== # Reference: https://kaldir.vc.in.tum.de/scannetpp/ # Note: High-quality indoor RGB-D dataset with iPhone and DSLR images. # Root directory for ScanNet++ evaluation data SCANNETPP_EVAL_DATA_ROOT = "workspace/benchmark_dataset/scannetpp" # List of ScanNet++ evaluation scenes SCANNETPP_SCENES = [ "09c1414f1b", "1ada7a0617", "40aec5fffa", "3e8bba0176", "acd95847c5", "578511c8a9", "5f99900f09", "c4c04e6d6c", "f3d64c30f8", "7bc286c1b6", "c5439f4607", "286b55a2bf", "fb5a96b1a2", "7831862f02", "38d58a7a31", "bde1e479ad", "9071e139d9", "21d970d8de", "bcd2436daf", "cc5237fd77", ] # Input resolution for ScanNet++ (after undistortion and resize) SCANNETPP_INPUT_H = 768 SCANNETPP_INPUT_W = 1024 # TSDF fusion hyperparameters (indoor scenes) SCANNETPP_VOXEL_LENGTH = 0.02 # Voxel size for TSDF (meters) SCANNETPP_SDF_TRUNC = 0.15 # SDF truncation distance (meters) SCANNETPP_MAX_DEPTH = 5.0 # Maximum depth for integration (meters) # Point cloud sampling SCANNETPP_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh # 3D reconstruction evaluation hyperparameters SCANNETPP_EVAL_THRESHOLD = 0.05 # Distance threshold for precision/recall (meters) SCANNETPP_DOWN_SAMPLE = 0.02 # Voxel size for evaluation downsampling (meters) # ============================================================================== # HiRoom Dataset Configuration # ============================================================================== # Note: Indoor RGB-D dataset. # Root directory for HiRoom evaluation data HIROOM_EVAL_DATA_ROOT = "workspace/benchmark_dataset/hiroom/data" HIROOM_GT_ROOT_PATH = "workspace/benchmark_dataset/hiroom/fused_pcd" HIROOM_SCENE_LIST_PATH = "workspace/benchmark_dataset/hiroom/selected_scene_list_val.txt" # TSDF fusion hyperparameters (indoor scenes) HIROOM_VOXEL_LENGTH = 4.0 / 512.0 # Voxel size for TSDF (meters) HIROOM_SDF_TRUNC = 0.04 # SDF truncation distance (meters) HIROOM_MAX_DEPTH = 10000.0 # Maximum depth for integration (no truncation) # Point cloud sampling HIROOM_SAMPLING_NUMBER = 1_000_000 # Number of points to sample from mesh # 3D reconstruction evaluation hyperparameters HIROOM_EVAL_THRESHOLD = 0.05 # Distance threshold for precision/recall (meters) HIROOM_DOWN_SAMPLE = 4.0 / 512.0 # Voxel size for evaluation downsampling (meters) # ============================================================================================================== # depth_anything_3/model/reference_view_selector.py # ============================================================================================================== # Copyright (c) 2025 ByteDance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Reference View Selection Strategies This module provides different strategies for selecting a reference view from multiple input views in multi-view depth estimation. """ import torch from typing import Literal RefViewStrategy = Literal["first", "middle", "saddle_balanced", "saddle_sim_range"] def select_reference_view( x: torch.Tensor, strategy: RefViewStrategy = "saddle_balanced", ) -> torch.Tensor: """ Select a reference view from multiple views using the specified strategy. Args: x: Input tensor of shape (B, S, N, C) where B = batch size S = number of views N = number of tokens C = channel dimension strategy: Selection strategy, one of: - "first": Always select the first view - "middle": Select the middle view - "saddle_balanced": Select view with balanced features across multiple metrics - "saddle_sim_range": Select view with largest similarity range Returns: b_idx: Tensor of shape (B,) containing the selected view index for each batch """ B, S, N, C = x.shape # For single view, no reordering needed if S <= 1: return torch.zeros(B, dtype=torch.long, device=x.device) # Simple position-based strategies if strategy == "first": return torch.zeros(B, dtype=torch.long, device=x.device) elif strategy == "middle": return torch.full((B,), S // 2, dtype=torch.long, device=x.device) # Feature-based strategies require normalized class tokens # Extract and normalize class tokens (first token of each view) img_class_feat = x[:, :, 0] / x[:, :, 0].norm(dim=-1, keepdim=True) # B S C if strategy == "saddle_balanced": # Select view with balanced features across multiple metrics # Compute similarity matrix sim = torch.matmul(img_class_feat, img_class_feat.transpose(1, 2)) # B S S sim_no_diag = sim - torch.eye(S, device=sim.device).unsqueeze(0) sim_score = sim_no_diag.sum(dim=-1) / (S - 1) # B S feat_norm = x[:, :, 0].norm(dim=-1) # B S feat_var = img_class_feat.var(dim=-1) # B S # Normalize all metrics to [0, 1] def normalize_metric(metric): min_val = metric.min(dim=1, keepdim=True).values max_val = metric.max(dim=1, keepdim=True).values return (metric - min_val) / (max_val - min_val + 1e-8) sim_score_norm = normalize_metric(sim_score) norm_norm = normalize_metric(feat_norm) var_norm = normalize_metric(feat_var) # Select view closest to the median (0.5) across all metrics balance_score = ( (sim_score_norm - 0.5).abs() + (norm_norm - 0.5).abs() + (var_norm - 0.5).abs() ) b_idx = balance_score.argmin(dim=1) elif strategy == "saddle_sim_range": # Select view with largest similarity range (max - min) sim = torch.matmul(img_class_feat, img_class_feat.transpose(1, 2)) # B S S sim_no_diag = sim - torch.eye(S, device=sim.device).unsqueeze(0) sim_max = sim_no_diag.max(dim=-1).values # B S sim_min = sim_no_diag.min(dim=-1).values # B S sim_range = sim_max - sim_min b_idx = sim_range.argmax(dim=1) else: raise ValueError( f"Unknown reference view selection strategy: {strategy}. " f"Must be one of: 'first', 'middle', 'saddle_balanced', 'saddle_sim_range'" ) return b_idx def reorder_by_reference( x: torch.Tensor, b_idx: torch.Tensor, ) -> torch.Tensor: """ Reorder views to place the selected reference view first. Args: x: Input tensor of shape (B, S, N, C) b_idx: Reference view indices of shape (B,) Returns: Reordered tensor with reference view at position 0 Example: If b_idx = [2] and S = 5 (views [0,1,2,3,4]), result order is [2,0,1,3,4] (ref_idx first, then others in order) """ B, S = x.shape[0], x.shape[1] # For single view, no reordering needed if S <= 1: return x # Create position indices: (B, S) where each row is [0, 1, 2, ..., S-1] positions = torch.arange(S, device=x.device).unsqueeze(0).expand(B, -1) # B S # For each position, determine which original index it should take # Position 0 gets ref_idx # Position 1 to ref_idx gets indices 0 to ref_idx-1 # Position ref_idx+1 to S-1 gets indices ref_idx+1 to S-1 b_idx_expanded = b_idx.unsqueeze(1) # B 1 # Create the reordering indices # For positions 1 to ref_idx: map to indices 0 to ref_idx-1 (shift by -1) # For positions > ref_idx: keep the same reorder_indices = positions.clone() reorder_indices = torch.where( (positions > 0) & (positions <= b_idx_expanded), positions - 1, positions ) # Set position 0 to ref_idx reorder_indices[:, 0] = b_idx # Gather using advanced indexing batch_indices = torch.arange(B, device=x.device).unsqueeze(1) # B 1 x_reordered = x[batch_indices, reorder_indices] return x_reordered def restore_original_order( x: torch.Tensor, b_idx: torch.Tensor, ) -> torch.Tensor: """ Restore original view order after processing. Args: x: Reordered tensor of shape (B, S, ...) b_idx: Original reference view indices of shape (B,) Returns: Tensor with original view order restored Example: If original order was [0, 1, 2, 3, 4] and b_idx=2, reordered becomes [2, 0, 1, 3, 4] (reference at position 0), restore should return [0, 1, 2, 3, 4] (original order). """ B, S = x.shape[0], x.shape[1] # For single view, no restoration needed if S <= 1: return x # Create target position indices: (B, S) where each row is [0, 1, 2, ..., S-1] target_positions = torch.arange(S, device=x.device).unsqueeze(0).expand(B, -1) # B S # For each target position, determine which current position it comes from # Target position 0 to ref_idx-1 <- Current position 1 to ref_idx (shift by +1) # Target position ref_idx <- Current position 0 # Target position ref_idx+1 to S-1 <- Current position ref_idx+1 to S-1 (no change) b_idx_expanded = b_idx.unsqueeze(1) # B 1 # Create the restore indices restore_indices = torch.where( target_positions < b_idx_expanded, target_positions + 1, # Positions before ref_idx come from current position + 1 target_positions # Positions after ref_idx stay the same ) # Target position = ref_idx comes from current position 0 # Use scatter to set specific positions restore_indices = torch.scatter( restore_indices, dim=1, index=b_idx_expanded, src=torch.zeros_like(b_idx_expanded) ) # Gather using advanced indexing batch_indices = torch.arange(B, device=x.device).unsqueeze(1) # B 1 x_restored = x[batch_indices, restore_indices] return x_restored # ============================================================================================================== # depth_anything_3/model/dinov2/layers/drop_path.py # ============================================================================================================== # 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. # References: # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py # https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py from torch import nn 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) # ============================================================================================================== # depth_anything_3/model/dinov2/layers/layer_scale.py # ============================================================================================================== # 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. # Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 # noqa: E501 from typing import Union import torch from torch import Tensor, nn class LayerScale(nn.Module): def __init__( self, dim: int, init_values: Union[float, Tensor] = 1e-5, inplace: bool = False, ) -> None: super().__init__() self.dim = dim self.inplace = inplace self.init_values = init_values 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 extra_repr(self) -> str: return f"{self.dim}, init_values={self.init_values}, inplace={self.inplace}" # ============================================================================================================== # depth_anything_3/model/dinov2/layers/mlp.py # ============================================================================================================== # 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. # References: # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py # https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py from typing import Callable, Optional from torch import Tensor, nn 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 # ============================================================================================================== # depth_anything_3/model/dinov2/layers/patch_embed.py # ============================================================================================================== # 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. # References: # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py # https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py from typing import Callable, Optional, Tuple, Union import torch.nn as nn from torch import Tensor def make_2tuple(x): if isinstance(x, tuple): assert len(x) == 2 return x assert isinstance(x, int) return (x, x) class PatchEmbed(nn.Module): """ 2D image to patch embedding: (B,C,H,W) -> (B,N,D) Args: img_size: Image size. patch_size: Patch token size. in_chans: Number of input image channels. embed_dim: Number of linear projection output channels. norm_layer: Normalization layer. """ def __init__( self, img_size: Union[int, Tuple[int, int]] = 224, patch_size: Union[int, Tuple[int, int]] = 16, in_chans: int = 3, embed_dim: int = 768, norm_layer: Optional[Callable] = None, flatten_embedding: bool = True, ) -> None: super().__init__() image_HW = make_2tuple(img_size) patch_HW = make_2tuple(patch_size) patch_grid_size = ( image_HW[0] // patch_HW[0], image_HW[1] // patch_HW[1], ) self.img_size = image_HW self.patch_size = patch_HW self.patches_resolution = patch_grid_size self.num_patches = patch_grid_size[0] * patch_grid_size[1] self.in_chans = in_chans self.embed_dim = embed_dim self.flatten_embedding = flatten_embedding self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW) self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() def forward(self, x: Tensor) -> Tensor: _, _, H, W = x.shape patch_H, patch_W = self.patch_size assert ( H % patch_H == 0 ), f"Input image height {H} is not a multiple of patch height {patch_H}" assert ( W % patch_W == 0 ), f"Input image width {W} is not a multiple of patch width: {patch_W}" x = self.proj(x) # B C H W H, W = x.size(2), x.size(3) x = x.flatten(2).transpose(1, 2) # B HW C x = self.norm(x) if not self.flatten_embedding: x = x.reshape(-1, H, W, self.embed_dim) # B H W C return x def flops(self) -> float: Ho, Wo = self.patches_resolution flops = ( Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) ) if self.norm is not None: flops += Ho * Wo * self.embed_dim return flops # ============================================================================================================== # depth_anything_3/model/dinov2/layers/swiglu_ffn.py # ============================================================================================================== # 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 Callable, Optional import torch.nn.functional as F from torch import Tensor, nn class SwiGLUFFN(nn.Module): def __init__( self, in_features: int, hidden_features: Optional[int] = None, out_features: Optional[int] = None, act_layer: Callable[..., nn.Module] = None, 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.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) self.w3 = nn.Linear(hidden_features, out_features, bias=bias) def forward(self, x: Tensor) -> Tensor: x12 = self.w12(x) x1, x2 = x12.chunk(2, dim=-1) hidden = F.silu(x1) * x2 return self.w3(hidden) try: from xformers.ops import SwiGLU XFORMERS_AVAILABLE = True except ImportError: SwiGLU = SwiGLUFFN XFORMERS_AVAILABLE = False class SwiGLUFFNFused(SwiGLU): def __init__( self, in_features: int, hidden_features: Optional[int] = None, out_features: Optional[int] = None, act_layer: Callable[..., nn.Module] = None, drop: float = 0.0, bias: bool = True, ) -> None: out_features = out_features or in_features hidden_features = hidden_features or in_features hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 super().__init__( in_features=in_features, hidden_features=hidden_features, out_features=out_features, bias=bias, ) # ============================================================================================================== # depth_anything_3/model/dinov2/layers/attention.py # ============================================================================================================== # 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. # References: # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py # https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py import logging import torch.nn.functional as F from torch import Tensor, nn _dinov2_logger = logging.getLogger("dinov2") 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, norm_layer: nn.Module = nn.LayerNorm, qk_norm: bool = False, fused_attn: bool = True, # use F.scaled_dot_product_attention or not rope=None, ) -> None: super().__init__() assert dim % num_heads == 0, "dim should be divisible by num_heads" self.num_heads = num_heads head_dim = dim // num_heads self.scale = head_dim**-0.5 self.fused_attn = fused_attn self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.q_norm = norm_layer(head_dim) if qk_norm else nn.Identity() self.k_norm = norm_layer(head_dim) if qk_norm else nn.Identity() self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim, bias=proj_bias) self.proj_drop = nn.Dropout(proj_drop) self.rope = rope def forward(self, x: Tensor, pos=None, attn_mask=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], qkv[1], qkv[2] q, k = self.q_norm(q), self.k_norm(k) if self.rope is not None and pos is not None: q = self.rope(q, pos) k = self.rope(k, pos) if self.fused_attn: x = F.scaled_dot_product_attention( q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0, attn_mask=( (attn_mask)[:, None].repeat(1, self.num_heads, 1, 1) if attn_mask is not None else None ), ) else: q = q * self.scale attn = q @ k.transpose(-2, -1) attn = attn.softmax(dim=-1) attn = self.attn_drop(attn) x = attn @ v x = x.transpose(1, 2).reshape(B, N, C) x = self.proj(x) x = self.proj_drop(x) return x 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 # ============================================================================================================== # depth_anything_3/model/dinov2/layers/block.py # ============================================================================================================== # flake8: noqa: F821 # 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. # References: # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py # https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py import logging from typing import Callable, Optional import torch from torch import Tensor, nn _dinov2_logger = logging.getLogger("dinov2") XFORMERS_AVAILABLE = True class Block(nn.Module): def __init__( self, dim: int, num_heads: int, mlp_ratio: float = 4.0, qkv_bias: bool = False, proj_bias: bool = True, ffn_bias: bool = True, drop: float = 0.0, attn_drop: float = 0.0, init_values=None, drop_path: float = 0.0, act_layer: Callable[..., nn.Module] = nn.GELU, norm_layer: Callable[..., nn.Module] = nn.LayerNorm, attn_class: Callable[..., nn.Module] = Attention, ffn_layer: Callable[..., nn.Module] = Mlp, qk_norm: bool = False, rope=None, ln_eps: float = 1e-6, ) -> None: super().__init__() # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}") self.norm1 = norm_layer(dim, eps=ln_eps) self.attn = attn_class( dim, num_heads=num_heads, qkv_bias=qkv_bias, proj_bias=proj_bias, attn_drop=attn_drop, proj_drop=drop, qk_norm=qk_norm, rope=rope, ) self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim, eps=ln_eps) mlp_hidden_dim = int(dim * mlp_ratio) self.mlp = ffn_layer( in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop, bias=ffn_bias, ) self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.sample_drop_ratio = drop_path def forward(self, x: Tensor, pos=None, attn_mask=None) -> Tensor: def attn_residual_func(x: Tensor, pos=None, attn_mask=None) -> Tensor: return self.ls1(self.attn(self.norm1(x), pos=pos, attn_mask=attn_mask)) def ffn_residual_func(x: Tensor) -> Tensor: return self.ls2(self.mlp(self.norm2(x))) if self.training and self.sample_drop_ratio > 0.1: # the overhead is compensated only for a drop path rate larger than 0.1 x = drop_add_residual_stochastic_depth( x, residual_func=attn_residual_func, sample_drop_ratio=self.sample_drop_ratio, pos=pos, ) x = drop_add_residual_stochastic_depth( x, residual_func=ffn_residual_func, sample_drop_ratio=self.sample_drop_ratio, ) elif self.training and self.sample_drop_ratio > 0.0: x = x + self.drop_path1(attn_residual_func(x, pos=pos, attn_mask=attn_mask)) x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2 else: x = x + attn_residual_func(x, pos=pos, attn_mask=attn_mask) x = x + ffn_residual_func(x) return x def drop_add_residual_stochastic_depth( x: Tensor, residual_func: Callable[[Tensor], Tensor], sample_drop_ratio: float = 0.0, pos: Optional[Tensor] = None, ) -> Tensor: # 1) extract subset using permutation b, n, d = x.shape sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) brange = (torch.randperm(b, device=x.device))[:sample_subset_size] x_subset = x[brange] # 2) apply residual_func to get residual if pos is not None: # if necessary, apply rope to the subset pos = pos[brange] residual = residual_func(x_subset, pos=pos) else: residual = residual_func(x_subset) x_flat = x.flatten(1) residual = residual.flatten(1) residual_scale_factor = b / sample_subset_size # 3) add the residual x_plus_residual = torch.index_add( x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor ) return x_plus_residual.view_as(x) def get_branges_scales(x, sample_drop_ratio=0.0): b, n, d = x.shape sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) brange = (torch.randperm(b, device=x.device))[:sample_subset_size] residual_scale_factor = b / sample_subset_size return brange, residual_scale_factor # ============================================================================================================== # depth_anything_3/model/dinov2/layers/rope.py # ============================================================================================================== # Copyright (c) Meta Platforms, Inc. and affiliates. # # This source code is licensed under the Apache License, Version 2.0 # found in the LICENSE file in the root directory of this source tree. # Implementation of 2D Rotary Position Embeddings (RoPE). # This module provides a clean implementation of 2D Rotary Position Embeddings, # which extends the original RoPE concept to handle 2D spatial positions. # Inspired by: # https://github.com/meta-llama/codellama/blob/main/llama/model.py # https://github.com/naver-ai/rope-vit from typing import Dict, Tuple import torch import torch.nn as nn import torch.nn.functional as F 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): """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]] = {} 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) # ============================================================================================================== # depth_anything_3/model/dinov2/vision_transformer.py # ============================================================================================================== # Copyright (c) Meta Platforms, Inc. and affiliates. # # This source code is licensed under the Apache License, Version 2.0 # found in the LICENSE file in the root directory of this source tree. # References: # https://github.com/facebookresearch/dino/blob/main/vision_transformer.py # https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py import math from typing import Callable, List, Sequence, Tuple, Union import numpy as np import torch import torch.nn as nn import torch.utils.checkpoint from einops import rearrange # logger = logging.getLogger("dinov2") def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): """ embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D) """ assert embed_dim % 2 == 0 omega = np.arange(embed_dim // 2, dtype=float) omega /= embed_dim / 2.0 omega = 1.0 / 10000**omega # (D/2,) pos = pos.reshape(-1) # (M,) out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product emb_sin = np.sin(out) # (M, D/2) emb_cos = np.cos(out) # (M, D/2) emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) return emb def named_apply( fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False ) -> nn.Module: if not depth_first and include_root: fn(module=module, name=name) for child_name, child_module in module.named_children(): child_name = ".".join((name, child_name)) if name else child_name named_apply( fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True ) if depth_first and include_root: fn(module=module, name=name) return module class BlockChunk(nn.ModuleList): def forward(self, x): for b in self: x = b(x) return x class DinoVisionTransformer(nn.Module): def __init__( self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True, ffn_bias=True, proj_bias=True, drop_path_rate=0.0, drop_path_uniform=False, init_values=1.0, # for layerscale: None or 0 => no layerscale embed_layer=PatchEmbed, act_layer=nn.GELU, block_fn=Block, ffn_layer="mlp", block_chunks=1, num_register_tokens=0, interpolate_antialias=False, interpolate_offset=0.1, alt_start=-1, qknorm_start=-1, rope_start=-1, rope_freq=100, plus_cam_token=False, cat_token=True, ): """ Args: img_size (int, tuple): input image size patch_size (int, tuple): patch size in_chans (int): number of input channels embed_dim (int): embedding dimension depth (int): depth of transformer num_heads (int): number of attention heads mlp_ratio (int): ratio of mlp hidden dim to embedding dim qkv_bias (bool): enable bias for qkv if True proj_bias (bool): enable bias for proj in attn if True ffn_bias (bool): enable bias for ffn if True weight_init (str): weight init scheme init_values (float): layer-scale init values embed_layer (nn.Module): patch embedding layer act_layer (nn.Module): MLP activation layer block_fn (nn.Module): transformer block class ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity" block_chunks: (int) split block sequence into block_chunks units for FSDP wrap num_register_tokens: (int) number of extra cls tokens (so-called "registers") interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings """ super().__init__() self.patch_start_idx = 1 norm_layer = nn.LayerNorm self.num_features = self.embed_dim = ( embed_dim # num_features for consistency with other models ) self.alt_start = alt_start self.qknorm_start = qknorm_start self.rope_start = rope_start self.cat_token = cat_token self.num_tokens = 1 self.n_blocks = depth self.num_heads = num_heads self.patch_size = patch_size self.num_register_tokens = num_register_tokens self.interpolate_antialias = interpolate_antialias self.interpolate_offset = interpolate_offset self.patch_embed = embed_layer( img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim ) num_patches = self.patch_embed.num_patches self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if self.alt_start != -1: self.camera_token = nn.Parameter(torch.randn(1, 2, embed_dim)) self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) assert num_register_tokens >= 0 self.register_tokens = ( nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None ) if drop_path_uniform is True: dpr = [drop_path_rate] * depth else: dpr = [ x.item() for x in torch.linspace(0, drop_path_rate, depth) ] # stochastic depth decay rule if ffn_layer == "mlp": logger.info("using MLP layer as FFN") ffn_layer = Mlp elif ffn_layer == "swiglufused" or ffn_layer == "swiglu": logger.info("using SwiGLU layer as FFN") ffn_layer = SwiGLUFFNFused elif ffn_layer == "identity": logger.info("using Identity layer as FFN") def f(*args, **kwargs): return nn.Identity() ffn_layer = f else: raise NotImplementedError if self.rope_start != -1: self.rope = RotaryPositionEmbedding2D(frequency=rope_freq) if rope_freq > 0 else None self.position_getter = PositionGetter() if self.rope is not None else None else: self.rope = None blocks_list = [ block_fn( dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, proj_bias=proj_bias, ffn_bias=ffn_bias, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer, ffn_layer=ffn_layer, init_values=init_values, qk_norm=i >= qknorm_start if qknorm_start != -1 else False, rope=self.rope if i >= rope_start and rope_start != -1 else None, ) for i in range(depth) ] self.blocks = nn.ModuleList(blocks_list) self.norm = norm_layer(embed_dim) def interpolate_pos_encoding(self, x, w, h): previous_dtype = x.dtype npatch = x.shape[1] - 1 N = self.pos_embed.shape[1] - 1 if npatch == N and w == h: return self.pos_embed pos_embed = self.pos_embed.float() class_pos_embed = pos_embed[:, 0] patch_pos_embed = pos_embed[:, 1:] dim = x.shape[-1] w0 = w // self.patch_size h0 = h // self.patch_size M = int(math.sqrt(N)) # Recover the number of patches in each dimension assert N == M * M kwargs = {} if self.interpolate_offset: # Historical kludge: add a small number to avoid floating point error in the # interpolation, see https://github.com/facebookresearch/dino/issues/8 # Note: still needed for backward-compatibility, the underlying operators are using # both output size and scale factors sx = float(w0 + self.interpolate_offset) / M sy = float(h0 + self.interpolate_offset) / M kwargs["scale_factor"] = (sx, sy) else: # Simply specify an output size instead of a scale factor kwargs["size"] = (w0, h0) patch_pos_embed = nn.functional.interpolate( patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2), mode="bicubic", antialias=self.interpolate_antialias, **kwargs, ) assert (w0, h0) == patch_pos_embed.shape[-2:] patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype) def prepare_cls_token(self, B, S): cls_token = self.cls_token.expand(B, S, -1) cls_token = cls_token.reshape(B * S, -1, self.embed_dim) return cls_token def prepare_tokens_with_masks(self, x, masks=None, cls_token=None, **kwargs): B, S, nc, w, h = x.shape x = rearrange(x, "b s c h w -> (b s) c h w") x = self.patch_embed(x) if masks is not None: x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x) cls_token = self.prepare_cls_token(B, S) x = torch.cat((cls_token, x), dim=1) x = x + self.interpolate_pos_encoding(x, w, h) if self.register_tokens is not None: x = torch.cat( ( x[:, :1], self.register_tokens.expand(x.shape[0], -1, -1), x[:, 1:], ), dim=1, ) x = rearrange(x, "(b s) n c -> b s n c", b=B, s=S) return x def _prepare_rope(self, B, S, H, W, device): pos = None pos_nodiff = None if self.rope is not None: pos = self.position_getter( B * S, H // self.patch_size, W // self.patch_size, device=device ) pos = rearrange(pos, "(b s) n c -> b s n c", b=B) pos_nodiff = torch.zeros_like(pos).to(pos.dtype) if self.patch_start_idx > 0: pos = pos + 1 pos_special = torch.zeros(B * S, self.patch_start_idx, 2).to(device).to(pos.dtype) pos_special = rearrange(pos_special, "(b s) n c -> b s n c", b=B) pos = torch.cat([pos_special, pos], dim=2) pos_nodiff = pos_nodiff + 1 pos_nodiff = torch.cat([pos_special, pos_nodiff], dim=2) return pos, pos_nodiff def _get_intermediate_layers_not_chunked(self, x, n=1, export_feat_layers=[], **kwargs): B, S, _, H, W = x.shape x = self.prepare_tokens_with_masks(x) output, total_block_len, aux_output = [], len(self.blocks), [] blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n pos, pos_nodiff = self._prepare_rope(B, S, H, W, x.device) for i, blk in enumerate(self.blocks): if i < self.rope_start or self.rope is None: g_pos, l_pos = None, None else: g_pos = pos_nodiff l_pos = pos if self.alt_start != -1 and (i == self.alt_start - 1) and x.shape[1] >= THRESH_FOR_REF_SELECTION and kwargs.get("cam_token", None) is None: # Select reference view using configured strategy strategy = kwargs.get("ref_view_strategy", "saddle_balanced") logger.info(f"Selecting reference view using strategy: {strategy}") b_idx = select_reference_view(x, strategy=strategy) # Reorder views to place reference view first x = reorder_by_reference(x, b_idx) local_x = reorder_by_reference(local_x, b_idx) if self.alt_start != -1 and i == self.alt_start: if kwargs.get("cam_token", None) is not None: logger.info("Using camera conditions provided by the user") cam_token = kwargs.get("cam_token") else: ref_token = self.camera_token[:, :1].expand(B, -1, -1) src_token = self.camera_token[:, 1:].expand(B, S - 1, -1) cam_token = torch.cat([ref_token, src_token], dim=1) x[:, :, 0] = cam_token if self.alt_start != -1 and i >= self.alt_start and i % 2 == 1: x = self.process_attention( x, blk, "global", pos=g_pos, attn_mask=kwargs.get("attn_mask", None) ) else: x = self.process_attention(x, blk, "local", pos=l_pos) local_x = x if i in blocks_to_take: out_x = torch.cat([local_x, x], dim=-1) if self.cat_token else x # Restore original view order if reordering was applied if x.shape[1] >= THRESH_FOR_REF_SELECTION and self.alt_start != -1 and 'b_idx' in locals(): out_x = restore_original_order(out_x, b_idx) output.append((out_x[:, :, 0], out_x)) if i in export_feat_layers: aux_output.append(x) return output, aux_output def process_attention(self, x, block, attn_type="global", pos=None, attn_mask=None): b, s, n = x.shape[:3] if attn_type == "local": x = rearrange(x, "b s n c -> (b s) n c") if pos is not None: pos = rearrange(pos, "b s n c -> (b s) n c") elif attn_type == "global": x = rearrange(x, "b s n c -> b (s n) c") if pos is not None: pos = rearrange(pos, "b s n c -> b (s n) c") else: raise ValueError(f"Invalid attention type: {attn_type}") x = block(x, pos=pos, attn_mask=attn_mask) if attn_type == "local": x = rearrange(x, "(b s) n c -> b s n c", b=b, s=s) elif attn_type == "global": x = rearrange(x, "b (s n) c -> b s n c", b=b, s=s) return x def get_intermediate_layers( self, x: torch.Tensor, n: Union[int, Sequence] = 1, # Layers or n last layers to take export_feat_layers: List[int] = [], **kwargs, ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]: outputs, aux_outputs = self._get_intermediate_layers_not_chunked( x, n, export_feat_layers=export_feat_layers, **kwargs ) camera_tokens = [out[0] for out in outputs] if outputs[0][1].shape[-1] == self.embed_dim: outputs = [self.norm(out[1]) for out in outputs] elif outputs[0][1].shape[-1] == (self.embed_dim * 2): outputs = [ torch.cat( [out[1][..., : self.embed_dim], self.norm(out[1][..., self.embed_dim :])], dim=-1, ) for out in outputs ] else: raise ValueError(f"Invalid output shape: {outputs[0][1].shape}") aux_outputs = [self.norm(out) for out in aux_outputs] outputs = [out[..., 1 + self.num_register_tokens :, :] for out in outputs] aux_outputs = [out[..., 1 + self.num_register_tokens :, :] for out in aux_outputs] return tuple(zip(outputs, camera_tokens)), aux_outputs def vit_small(patch_size=16, num_register_tokens=0, depth=12, **kwargs): model = DinoVisionTransformer( patch_size=patch_size, embed_dim=384, depth=depth, num_heads=6, mlp_ratio=4, # block_fn=partial(Block, attn_class=MemEffAttention), num_register_tokens=num_register_tokens, **kwargs, ) return model def vit_base(patch_size=16, num_register_tokens=0, depth=12, **kwargs): model = DinoVisionTransformer( patch_size=patch_size, embed_dim=768, depth=depth, num_heads=12, mlp_ratio=4, # block_fn=partial(Block, attn_class=MemEffAttention), num_register_tokens=num_register_tokens, **kwargs, ) return model def vit_large(patch_size=16, num_register_tokens=0, depth=24, **kwargs): model = DinoVisionTransformer( patch_size=patch_size, embed_dim=1024, depth=depth, num_heads=16, mlp_ratio=4, # block_fn=partial(Block, attn_class=MemEffAttention), num_register_tokens=num_register_tokens, **kwargs, ) return model def vit_giant2(patch_size=16, num_register_tokens=0, depth=40, **kwargs): """ Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64 """ model = DinoVisionTransformer( patch_size=patch_size, embed_dim=1536, depth=depth, num_heads=24, mlp_ratio=4, num_register_tokens=num_register_tokens, **kwargs, ) return model # ============================================================================================================== # depth_anything_3/model/dinov2/dinov2.py # ============================================================================================================== # Copyright (c) Meta Platforms, Inc. and affiliates. # # This source code is licensed under the Apache License, Version 2.0 # found in the LICENSE file in the root directory of this source tree. # References: # https://github.com/facebookresearch/dino/blob/main/vision_transformer.py # https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py from typing import List import torch.nn as nn class DinoV2(nn.Module): def __init__( self, name: str, out_layers: List[int], alt_start: int = -1, qknorm_start: int = -1, rope_start: int = -1, cat_token: bool = True, **kwargs, ): super().__init__() assert name in {"vits", "vitb", "vitl", "vitg"} self.name = name self.out_layers = out_layers self.alt_start = alt_start self.qknorm_start = qknorm_start self.rope_start = rope_start self.cat_token = cat_token encoder_map = { "vits": vit_small, "vitb": vit_base, "vitl": vit_large, "vitg": vit_giant2, } encoder_fn = encoder_map[self.name] ffn_layer = "swiglufused" if self.name == "vitg" else "mlp" self.pretrained = encoder_fn( img_size=518, patch_size=14, ffn_layer=ffn_layer, alt_start=alt_start, qknorm_start=qknorm_start, rope_start=rope_start, cat_token=cat_token, ) def forward(self, x, **kwargs): return self.pretrained.get_intermediate_layers( x, self.out_layers, **kwargs, )