Delete External Models to prevent HF tags

mapanything/models/external/README.md

@@ -1,5 +0,0 @@
-# External Model Code for Benchmarking & Re-Training
-
-This directory contains external model code that we use to train and benchmark external models fairly. These libraries are not part of the core MapAnything codebase and are included for only benchmarking purposes. The code in this directory is licensed under the same license as the source code from which it was derived, unless otherwise specified.
-
-The open-source Apache 2.0 License of MapAnything does not apply to these libraries.

mapanything/models/external/anycalib/__init__.py

@@ -1,95 +0,0 @@
-"""
-Inference wrapper for AnyCalib
-"""
-
-import torch
-from anycalib import AnyCalib
-
-from mapanything.utils.geometry import get_rays_in_camera_frame
-
-
-class AnyCalibWrapper(torch.nn.Module):
-    def __init__(
-        self,
-        name,
-        model_id="anycalib_pinhole",
-        **kwargs,
-    ):
-        super().__init__()
-        self.name = name
-        self.model_id = model_id
-
-        # Initialize the model
-        self.model = AnyCalib(model_id=self.model_id)
-
-    def forward(self, views):
-        """
-        Forward pass wrapper for AnyCalib.
-
-        Assumption:
-        - The number of input views is 1.
-        - The output camera model is pinhole (fx, fy, cx, cy).
-          This can be relaxed by not hardcoding the cam_id.
-
-        Args:
-            views (List[dict]): List of dictionaries containing the input views' images and instance information.
-                                Length of the list should be 1.
-                                Each dictionary should contain the following keys:
-                                    "img" (tensor): Image tensor of shape (B, C, H, W).
-                                    "data_norm_type" (list): ["identity"]
-
-        Returns:
-            List[dict]: A list containing the final outputs for the single view. Length of the list will be 1.
-        """
-        # Check that the number of input views is 1
-        assert len(views) == 1, "AnyCalib only supports 1 input view."
-
-        # Get input shape of the images and batch size per view
-        _, _, height, width = views[0]["img"].shape
-
-        # Check the data norm type
-        # AnyCalib expects a normalized image but without the DINOv2 mean and std applied ("identity")
-        data_norm_type = views[0]["data_norm_type"][0]
-        assert data_norm_type == "identity", (
-            "AnyCalib expects a normalized image but without the DINOv2 mean and std applied"
-        )
-
-        # Run AnyCalib inference
-        # Corresponding batched output dictionary:
-        # {
-        #      "intrinsics": List[(D_i,) tensors] for each camera model "i" at the original input resolution,
-        #      "fov_field": (B, N, 2) tensor with the regressed FoV field by the network. N≈320^2 (resolution close to the one seen during training),
-        #      "tangent_coords": alias for "fov_field",
-        #      "rays": (B, N, 3) tensor with the corresponding (via the exponential map) ray directions in the camera frame (x right, y down, z forward),
-        #      "pred_size": (H, W) tuple with the image size used by the network. It can be used e.g. for resizing the FoV/ray fields to the original image size.
-        # }
-        # For "pinhole" camera model, the intrinsics are (fx, fy, cx, cy).
-        model_outputs = self.model.predict(views[0]["img"], cam_id="pinhole")
-
-        # Convert the list of intrinsics to a tensor
-        intrinsics = []
-        for intrinsics_per_sample in model_outputs["intrinsics"]:
-            pred_fx, pred_fy, pred_cx, pred_cy = intrinsics_per_sample
-            intrinsics_per_sample = torch.tensor(
-                [
-                    [pred_fx, 0, pred_cx],
-                    [0, pred_fy, pred_cy],
-                    [0, 0, 1],
-                ],
-                device=views[0]["img"].device,
-            )
-            intrinsics.append(intrinsics_per_sample)
-
-        # Convert the list of intrinsics to a tensor of size (batch_size_per_view, 3, 3)
-        intrinsics = torch.stack(intrinsics)
-
-        # Get the ray directions
-        with torch.autocast("cuda", enabled=False):
-            _, ray_directions = get_rays_in_camera_frame(
-                intrinsics, height, width, normalize_to_unit_sphere=True
-            )
-
-        # Return the output in MapAnything format
-        res = [{"ray_directions": ray_directions, "intrinsics": intrinsics}]
-
-        return res

mapanything/models/external/dinov2/__init__.py

@@ -1,6 +0,0 @@
-# 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.
-
-__version__ = "0.0.1"

mapanything/models/external/dinov2/hub/__init__.py

@@ -1,4 +0,0 @@
-# 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.

mapanything/models/external/dinov2/hub/backbones.py

@@ -1,183 +0,0 @@
-# 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.
-
-from enum import Enum
-from typing import Union
-
-import torch
-
-from mapanything.models.external.dinov2.hub.utils import (
-    _DINOV2_BASE_URL,
-    _make_dinov2_model_name,
-)
-
-
-class Weights(Enum):
-    LVD142M = "LVD142M"
-
-
-def _make_dinov2_model(
-    *,
-    arch_name: str = "vit_large",
-    img_size: int = 518,
-    patch_size: int = 14,
-    init_values: float = 1.0,
-    ffn_layer: str = "mlp",
-    block_chunks: int = 0,
-    num_register_tokens: int = 0,
-    interpolate_antialias: bool = False,
-    interpolate_offset: float = 0.1,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.LVD142M,
-    **kwargs,
-):
-    from ..models import vision_transformer as vits
-
-    if isinstance(weights, str):
-        try:
-            weights = Weights[weights]
-        except KeyError:
-            raise AssertionError(f"Unsupported weights: {weights}")
-
-    model_base_name = _make_dinov2_model_name(arch_name, patch_size)
-    vit_kwargs = dict(
-        img_size=img_size,
-        patch_size=patch_size,
-        init_values=init_values,
-        ffn_layer=ffn_layer,
-        block_chunks=block_chunks,
-        num_register_tokens=num_register_tokens,
-        interpolate_antialias=interpolate_antialias,
-        interpolate_offset=interpolate_offset,
-    )
-    vit_kwargs.update(**kwargs)
-    model = vits.__dict__[arch_name](**vit_kwargs)
-
-    if pretrained:
-        model_full_name = _make_dinov2_model_name(
-            arch_name, patch_size, num_register_tokens
-        )
-        url = _DINOV2_BASE_URL + f"/{model_base_name}/{model_full_name}_pretrain.pth"
-        state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu")
-        model.load_state_dict(state_dict, strict=True)
-
-    return model
-
-
-def dinov2_vits14(
-    *, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs
-):
-    """
-    DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_small", pretrained=pretrained, weights=weights, **kwargs
-    )
-
-
-def dinov2_vitb14(
-    *, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs
-):
-    """
-    DINOv2 ViT-B/14 model (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_base", pretrained=pretrained, weights=weights, **kwargs
-    )
-
-
-def dinov2_vitl14(
-    *, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs
-):
-    """
-    DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_large", pretrained=pretrained, weights=weights, **kwargs
-    )
-
-
-def dinov2_vitg14(
-    *, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs
-):
-    """
-    DINOv2 ViT-g/14 model (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_giant2",
-        ffn_layer="swiglufused",
-        weights=weights,
-        pretrained=pretrained,
-        **kwargs,
-    )
-
-
-def dinov2_vits14_reg(
-    *, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs
-):
-    """
-    DINOv2 ViT-S/14 model with registers (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_small",
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )
-
-
-def dinov2_vitb14_reg(
-    *, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs
-):
-    """
-    DINOv2 ViT-B/14 model with registers (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_base",
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )
-
-
-def dinov2_vitl14_reg(
-    *, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs
-):
-    """
-    DINOv2 ViT-L/14 model with registers (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_large",
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )
-
-
-def dinov2_vitg14_reg(
-    *, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs
-):
-    """
-    DINOv2 ViT-g/14 model with registers (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_giant2",
-        ffn_layer="swiglufused",
-        weights=weights,
-        pretrained=pretrained,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )

mapanything/models/external/dinov2/hub/utils.py

@@ -1,42 +0,0 @@
-# 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.
-
-import itertools
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-_DINOV2_BASE_URL = "https://dl.fbaipublicfiles.com/dinov2"
-
-
-def _make_dinov2_model_name(
-    arch_name: str, patch_size: int, num_register_tokens: int = 0
-) -> str:
-    compact_arch_name = arch_name.replace("_", "")[:4]
-    registers_suffix = f"_reg{num_register_tokens}" if num_register_tokens else ""
-    return f"dinov2_{compact_arch_name}{patch_size}{registers_suffix}"
-
-
-class CenterPadding(nn.Module):
-    def __init__(self, multiple):
-        super().__init__()
-        self.multiple = multiple
-
-    def _get_pad(self, size):
-        new_size = math.ceil(size / self.multiple) * self.multiple
-        pad_size = new_size - size
-        pad_size_left = pad_size // 2
-        pad_size_right = pad_size - pad_size_left
-        return pad_size_left, pad_size_right
-
-    @torch.inference_mode()
-    def forward(self, x):
-        pads = list(
-            itertools.chain.from_iterable(self._get_pad(m) for m in x.shape[:1:-1])
-        )
-        output = F.pad(x, pads)
-        return output

mapanything/models/external/dinov2/layers/__init__.py

@@ -1,14 +0,0 @@
-# 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.
-
-from mapanything.models.external.dinov2.layers.dino_head import DINOHead  # noqa
-from mapanything.models.external.dinov2.layers.mlp import Mlp  # noqa
-from mapanything.models.external.dinov2.layers.patch_embed import PatchEmbed  # noqa
-from mapanything.models.external.dinov2.layers.swiglu_ffn import (
-    SwiGLUFFN,  # noqa
-    SwiGLUFFNFused,  # noqa
-)
-from mapanything.models.external.dinov2.layers.block import NestedTensorBlock  # noqa
-from mapanything.models.external.dinov2.layers.attention import MemEffAttention  # noqa

mapanything/models/external/dinov2/layers/attention.py

@@ -1,90 +0,0 @@
-# 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/master/vision_transformer.py
-#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
-
-import logging
-import os
-
-from torch import nn, Tensor
-
-logger = logging.getLogger("dinov2")
-
-
-XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
-try:
-    if XFORMERS_ENABLED:
-        from xformers.ops import memory_efficient_attention, unbind
-
-        XFORMERS_AVAILABLE = True
-        # warnings.warn("xFormers is available (Attention)")
-    else:
-        # warnings.warn("xFormers is disabled (Attention)")
-        raise ImportError
-except ImportError:
-    XFORMERS_AVAILABLE = False
-    # warnings.warn("xFormers is not available (Attention)")
-
-
-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, attn_bias=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]
-        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, attn_bias=None) -> Tensor:
-        if not XFORMERS_AVAILABLE:
-            if attn_bias is not None:
-                raise AssertionError("xFormers is required for using nested tensors")
-            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)
-
-        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
-        x = x.reshape([B, N, C])
-
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x

mapanything/models/external/dinov2/layers/block.py

@@ -1,290 +0,0 @@
-# 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/master/vision_transformer.py
-#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
-
-import logging
-import os
-from typing import Any, Callable, Dict, List, Tuple
-
-import torch
-from torch import nn, Tensor
-
-from mapanything.models.external.dinov2.layers.attention import (
-    Attention,
-    MemEffAttention,
-)
-from mapanything.models.external.dinov2.layers.drop_path import DropPath
-from mapanything.models.external.dinov2.layers.layer_scale import LayerScale
-from mapanything.models.external.dinov2.layers.mlp import Mlp
-
-logger = logging.getLogger("dinov2")
-
-
-XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
-try:
-    if XFORMERS_ENABLED:
-        from xformers.ops import fmha, index_select_cat, scaled_index_add
-
-        XFORMERS_AVAILABLE = True
-        # warnings.warn("xFormers is available (Block)")
-    else:
-        # warnings.warn("xFormers is disabled (Block)")
-        raise ImportError
-except ImportError:
-    XFORMERS_AVAILABLE = False
-    # warnings.warn("xFormers is not available (Block)")
-
-
-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,
-    ) -> None:
-        super().__init__()
-        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
-        self.norm1 = norm_layer(dim)
-        self.attn = attn_class(
-            dim,
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            proj_bias=proj_bias,
-            attn_drop=attn_drop,
-            proj_drop=drop,
-        )
-        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)
-        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) -> Tensor:
-        def attn_residual_func(x: Tensor) -> Tensor:
-            return self.ls1(self.attn(self.norm1(x)))
-
-        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,
-            )
-            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))
-            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
-        else:
-            x = x + attn_residual_func(x)
-            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,
-) -> 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
-    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
-
-
-def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
-    if scaling_vector is None:
-        x_flat = x.flatten(1)
-        residual = residual.flatten(1)
-        x_plus_residual = torch.index_add(
-            x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor
-        )
-    else:
-        x_plus_residual = scaled_index_add(
-            x,
-            brange,
-            residual.to(dtype=x.dtype),
-            scaling=scaling_vector,
-            alpha=residual_scale_factor,
-        )
-    return x_plus_residual
-
-
-attn_bias_cache: Dict[Tuple, Any] = {}
-
-
-def get_attn_bias_and_cat(x_list, branges=None):
-    """
-    this will perform the index select, cat the tensors, and provide the attn_bias from cache
-    """
-    batch_sizes = (
-        [b.shape[0] for b in branges]
-        if branges is not None
-        else [x.shape[0] for x in x_list]
-    )
-    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
-    if all_shapes not in attn_bias_cache.keys():
-        seqlens = []
-        for b, x in zip(batch_sizes, x_list):
-            for _ in range(b):
-                seqlens.append(x.shape[1])
-        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
-        attn_bias._batch_sizes = batch_sizes
-        attn_bias_cache[all_shapes] = attn_bias
-
-    if branges is not None:
-        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(
-            1, -1, x_list[0].shape[-1]
-        )
-    else:
-        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
-        cat_tensors = torch.cat(tensors_bs1, dim=1)
-
-    return attn_bias_cache[all_shapes], cat_tensors
-
-
-def drop_add_residual_stochastic_depth_list(
-    x_list: List[Tensor],
-    residual_func: Callable[[Tensor, Any], Tensor],
-    sample_drop_ratio: float = 0.0,
-    scaling_vector=None,
-) -> Tensor:
-    # 1) generate random set of indices for dropping samples in the batch
-    branges_scales = [
-        get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list
-    ]
-    branges = [s[0] for s in branges_scales]
-    residual_scale_factors = [s[1] for s in branges_scales]
-
-    # 2) get attention bias and index+concat the tensors
-    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
-
-    # 3) apply residual_func to get residual, and split the result
-    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
-
-    outputs = []
-    for x, brange, residual, residual_scale_factor in zip(
-        x_list, branges, residual_list, residual_scale_factors
-    ):
-        outputs.append(
-            add_residual(
-                x, brange, residual, residual_scale_factor, scaling_vector
-            ).view_as(x)
-        )
-    return outputs
-
-
-class NestedTensorBlock(Block):
-    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
-        """
-        x_list contains a list of tensors to nest together and run
-        """
-        assert isinstance(self.attn, MemEffAttention)
-
-        if self.training and self.sample_drop_ratio > 0.0:
-
-            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.attn(self.norm1(x), attn_bias=attn_bias)
-
-            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.mlp(self.norm2(x))
-
-            x_list = drop_add_residual_stochastic_depth_list(
-                x_list,
-                residual_func=attn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-                scaling_vector=self.ls1.gamma
-                if isinstance(self.ls1, LayerScale)
-                else None,
-            )
-            x_list = drop_add_residual_stochastic_depth_list(
-                x_list,
-                residual_func=ffn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-                scaling_vector=self.ls2.gamma
-                if isinstance(self.ls1, LayerScale)
-                else None,
-            )
-            return x_list
-        else:
-
-            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
-
-            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.ls2(self.mlp(self.norm2(x)))
-
-            attn_bias, x = get_attn_bias_and_cat(x_list)
-            x = x + attn_residual_func(x, attn_bias=attn_bias)
-            x = x + ffn_residual_func(x)
-            return attn_bias.split(x)
-
-    def forward(self, x_or_x_list):
-        if isinstance(x_or_x_list, Tensor):
-            return super().forward(x_or_x_list)
-        elif isinstance(x_or_x_list, list):
-            if not XFORMERS_AVAILABLE:
-                raise AssertionError("xFormers is required for using nested tensors")
-            return self.forward_nested(x_or_x_list)
-        else:
-            raise AssertionError

mapanything/models/external/dinov2/layers/dino_head.py

@@ -1,67 +0,0 @@
-# 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.
-
-import torch
-import torch.nn as nn
-from torch.nn.init import trunc_normal_
-from torch.nn.utils import weight_norm
-
-
-class DINOHead(nn.Module):
-    def __init__(
-        self,
-        in_dim,
-        out_dim,
-        use_bn=False,
-        nlayers=3,
-        hidden_dim=2048,
-        bottleneck_dim=256,
-        mlp_bias=True,
-    ):
-        super().__init__()
-        nlayers = max(nlayers, 1)
-        self.mlp = _build_mlp(
-            nlayers,
-            in_dim,
-            bottleneck_dim,
-            hidden_dim=hidden_dim,
-            use_bn=use_bn,
-            bias=mlp_bias,
-        )
-        self.apply(self._init_weights)
-        self.last_layer = weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
-        self.last_layer.weight_g.data.fill_(1)
-
-    def _init_weights(self, m):
-        if isinstance(m, nn.Linear):
-            trunc_normal_(m.weight, std=0.02)
-            if isinstance(m, nn.Linear) and m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-
-    def forward(self, x):
-        x = self.mlp(x)
-        eps = 1e-6 if x.dtype == torch.float16 else 1e-12
-        x = nn.functional.normalize(x, dim=-1, p=2, eps=eps)
-        x = self.last_layer(x)
-        return x
-
-
-def _build_mlp(
-    nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True
-):
-    if nlayers == 1:
-        return nn.Linear(in_dim, bottleneck_dim, bias=bias)
-    else:
-        layers = [nn.Linear(in_dim, hidden_dim, bias=bias)]
-        if use_bn:
-            layers.append(nn.BatchNorm1d(hidden_dim))
-        layers.append(nn.GELU())
-        for _ in range(nlayers - 2):
-            layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias))
-            if use_bn:
-                layers.append(nn.BatchNorm1d(hidden_dim))
-            layers.append(nn.GELU())
-        layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias))
-        return nn.Sequential(*layers)

mapanything/models/external/dinov2/layers/drop_path.py

@@ -1,36 +0,0 @@
-# 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/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(DropPath, self).__init__()
-        self.drop_prob = drop_prob
-
-    def forward(self, x):
-        return drop_path(x, self.drop_prob, self.training)

mapanything/models/external/dinov2/layers/layer_scale.py

@@ -1,26 +0,0 @@
-# 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.
-
-# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
-
-from typing import Union
-
-import torch
-from torch import nn, Tensor
-
-
-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

mapanything/models/external/dinov2/layers/mlp.py

@@ -1,40 +0,0 @@
-# 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/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 nn, Tensor
-
-
-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

mapanything/models/external/dinov2/layers/patch_embed.py

@@ -1,100 +0,0 @@
-# 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/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

mapanything/models/external/dinov2/layers/swiglu_ffn.py

@@ -1,71 +0,0 @@
-# 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.
-
-import os
-from typing import Callable, Optional
-
-import torch.nn.functional as F
-from torch import nn, Tensor
-
-
-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)
-
-
-XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
-try:
-    if XFORMERS_ENABLED:
-        from xformers.ops import SwiGLU
-
-        XFORMERS_AVAILABLE = True
-        # warnings.warn("xFormers is available (SwiGLU)")
-    else:
-        # warnings.warn("xFormers is disabled (SwiGLU)")
-        raise ImportError
-except ImportError:
-    SwiGLU = SwiGLUFFN
-    XFORMERS_AVAILABLE = False
-
-    # warnings.warn("xFormers is not available (SwiGLU)")
-
-
-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,
-        )

mapanything/models/external/dinov2/models/__init__.py

@@ -1,44 +0,0 @@
-# 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.
-
-import logging
-
-import mapanything.models.external.dinov2.models.vision_transformer as vits
-
-logger = logging.getLogger("dinov2")
-
-
-def build_model(args, only_teacher=False, img_size=224):
-    args.arch = args.arch.removesuffix("_memeff")
-    if "vit" in args.arch:
-        vit_kwargs = dict(
-            img_size=img_size,
-            patch_size=args.patch_size,
-            init_values=args.layerscale,
-            ffn_layer=args.ffn_layer,
-            block_chunks=args.block_chunks,
-            qkv_bias=args.qkv_bias,
-            proj_bias=args.proj_bias,
-            ffn_bias=args.ffn_bias,
-            num_register_tokens=args.num_register_tokens,
-            interpolate_offset=args.interpolate_offset,
-            interpolate_antialias=args.interpolate_antialias,
-        )
-        teacher = vits.__dict__[args.arch](**vit_kwargs)
-        if only_teacher:
-            return teacher, teacher.embed_dim
-        student = vits.__dict__[args.arch](
-            **vit_kwargs,
-            drop_path_rate=args.drop_path_rate,
-            drop_path_uniform=args.drop_path_uniform,
-        )
-        embed_dim = student.embed_dim
-    return student, teacher, embed_dim
-
-
-def build_model_from_cfg(cfg, only_teacher=False):
-    return build_model(
-        cfg.student, only_teacher=only_teacher, img_size=cfg.crops.global_crops_size
-    )

mapanything/models/external/dinov2/models/vision_transformer.py

@@ -1,448 +0,0 @@
-# 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 functools import partial
-from typing import Callable, Sequence, Tuple, Union
-
-import torch
-import torch.nn as nn
-from torch.nn.init import trunc_normal_
-from torch.utils.checkpoint import checkpoint
-
-from mapanything.models.external.dinov2.layers import (
-    MemEffAttention,
-    Mlp,
-    NestedTensorBlock as Block,
-    PatchEmbed,
-    SwiGLUFFNFused,
-)
-from mapanything.models.external.pi3.layers.attention import FlashAttention
-
-# logger = logging.getLogger("dinov2")
-
-
-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=None,  # 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,
-    ):
-        """
-        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
-            drop_path_rate (float): stochastic depth rate
-            drop_path_uniform (bool): apply uniform drop rate across blocks
-            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__()
-        norm_layer = partial(nn.LayerNorm, eps=1e-6)
-
-        self.num_features = self.embed_dim = (
-            embed_dim  # num_features for consistency with other models
-        )
-        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))
-        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
-
-        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,
-                attn_class=FlashAttention,
-            )
-            for i in range(depth)
-        ]
-        if block_chunks > 0:
-            self.chunked_blocks = True
-            chunked_blocks = []
-            chunksize = depth // block_chunks
-            for i in range(0, depth, chunksize):
-                # this is to keep the block index consistent if we chunk the block list
-                chunked_blocks.append(
-                    [nn.Identity()] * i + blocks_list[i : i + chunksize]
-                )
-            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
-        else:
-            self.chunked_blocks = False
-            self.blocks = nn.ModuleList(blocks_list)
-
-        self.norm = norm_layer(embed_dim)
-        self.head = nn.Identity()
-
-        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
-
-        self.init_weights()
-
-    def init_weights(self):
-        trunc_normal_(self.pos_embed, std=0.02)
-        nn.init.normal_(self.cls_token, std=1e-6)
-        if self.register_tokens is not None:
-            nn.init.normal_(self.register_tokens, std=1e-6)
-        named_apply(init_weights_vit_timm, self)
-
-    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_tokens_with_masks(self, x, masks=None):
-        B, nc, w, h = x.shape
-        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
-            )
-
-        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), 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,
-            )
-
-        return x
-
-    def forward_features_list(self, x_list, masks_list):
-        x = [
-            self.prepare_tokens_with_masks(x, masks)
-            for x, masks in zip(x_list, masks_list)
-        ]
-        for blk in self.blocks:
-            if self.training:
-                x = checkpoint(blk, x, use_reentrant=False)
-            else:
-                x = blk(x)
-
-        all_x = x
-        output = []
-        for x, masks in zip(all_x, masks_list):
-            x_norm = self.norm(x)
-            output.append(
-                {
-                    "x_norm_clstoken": x_norm[:, 0],
-                    "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
-                    "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
-                    "x_prenorm": x,
-                    "masks": masks,
-                }
-            )
-        return output
-
-    def forward_features(self, x, masks=None):
-        if isinstance(x, list):
-            return self.forward_features_list(x, masks)
-
-        x = self.prepare_tokens_with_masks(x, masks)
-
-        for blk in self.blocks:
-            if self.training:
-                x = checkpoint(blk, x, use_reentrant=False)
-            else:
-                x = blk(x)
-
-        x_norm = self.norm(x)
-        return {
-            "x_norm_clstoken": x_norm[:, 0],
-            "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
-            "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
-            "x_prenorm": x,
-            "masks": masks,
-        }
-
-    def _get_intermediate_layers_not_chunked(self, x, n=1):
-        x = self.prepare_tokens_with_masks(x)
-        # If n is an int, take the n last blocks. If it's a list, take them
-        output, total_block_len = [], len(self.blocks)
-        blocks_to_take = (
-            range(total_block_len - n, total_block_len) if isinstance(n, int) else n
-        )
-        for i, blk in enumerate(self.blocks):
-            x = blk(x)
-            if i in blocks_to_take:
-                output.append(x)
-        assert len(output) == len(blocks_to_take), (
-            f"only {len(output)} / {len(blocks_to_take)} blocks found"
-        )
-        return output
-
-    def _get_intermediate_layers_chunked(self, x, n=1):
-        x = self.prepare_tokens_with_masks(x)
-        output, i, total_block_len = [], 0, len(self.blocks[-1])
-        # If n is an int, take the n last blocks. If it's a list, take them
-        blocks_to_take = (
-            range(total_block_len - n, total_block_len) if isinstance(n, int) else n
-        )
-        for block_chunk in self.blocks:
-            for blk in block_chunk[i:]:  # Passing the nn.Identity()
-                x = blk(x)
-                if i in blocks_to_take:
-                    output.append(x)
-                i += 1
-        assert len(output) == len(blocks_to_take), (
-            f"only {len(output)} / {len(blocks_to_take)} blocks found"
-        )
-        return output
-
-    def get_intermediate_layers(
-        self,
-        x: torch.Tensor,
-        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
-        reshape: bool = False,
-        return_class_token: bool = False,
-        norm=True,
-    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
-        if self.chunked_blocks:
-            outputs = self._get_intermediate_layers_chunked(x, n)
-        else:
-            outputs = self._get_intermediate_layers_not_chunked(x, n)
-        if norm:
-            outputs = [self.norm(out) for out in outputs]
-        class_tokens = [out[:, 0] for out in outputs]
-        outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs]
-        if reshape:
-            B, _, w, h = x.shape
-            outputs = [
-                out.reshape(B, w // self.patch_size, h // self.patch_size, -1)
-                .permute(0, 3, 1, 2)
-                .contiguous()
-                for out in outputs
-            ]
-        if return_class_token:
-            return tuple(zip(outputs, class_tokens))
-        return tuple(outputs)
-
-    def forward(self, *args, is_training=False, **kwargs):
-        ret = self.forward_features(*args, **kwargs)
-        if is_training:
-            return ret
-        else:
-            return self.head(ret["x_norm_clstoken"])
-
-
-def init_weights_vit_timm(module: nn.Module, name: str = ""):
-    """ViT weight initialization, original timm impl (for reproducibility)"""
-    if isinstance(module, nn.Linear):
-        trunc_normal_(module.weight, std=0.02)
-        if module.bias is not None:
-            nn.init.zeros_(module.bias)
-
-
-def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=384,
-        depth=12,
-        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, **kwargs):
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=768,
-        depth=12,
-        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, **kwargs):
-    model = DinoVisionTransformer(
-        patch_size=patch_size,
-        embed_dim=1024,
-        depth=24,
-        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, **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=40,
-        num_heads=24,
-        mlp_ratio=4,
-        block_fn=partial(Block, attn_class=MemEffAttention),
-        num_register_tokens=num_register_tokens,
-        **kwargs,
-    )
-    return model

mapanything/models/external/dinov2/utils/__init__.py

@@ -1,4 +0,0 @@
-# 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.

mapanything/models/external/dinov2/utils/cluster.py

@@ -1,102 +0,0 @@
-# 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.
-
-import os
-from enum import Enum
-from pathlib import Path
-from typing import Any, Dict, Optional
-
-
-class ClusterType(Enum):
-    AWS = "aws"
-    FAIR = "fair"
-    RSC = "rsc"
-
-
-def _guess_cluster_type() -> ClusterType:
-    uname = os.uname()
-    if uname.sysname == "Linux":
-        if uname.release.endswith("-aws"):
-            # Linux kernel versions on AWS instances are of the form "5.4.0-1051-aws"
-            return ClusterType.AWS
-        elif uname.nodename.startswith("rsc"):
-            # Linux kernel versions on RSC instances are standard ones but hostnames start with "rsc"
-            return ClusterType.RSC
-
-    return ClusterType.FAIR
-
-
-def get_cluster_type(
-    cluster_type: Optional[ClusterType] = None,
-) -> Optional[ClusterType]:
-    if cluster_type is None:
-        return _guess_cluster_type()
-
-    return cluster_type
-
-
-def get_checkpoint_path(cluster_type: Optional[ClusterType] = None) -> Optional[Path]:
-    cluster_type = get_cluster_type(cluster_type)
-    if cluster_type is None:
-        return None
-
-    CHECKPOINT_DIRNAMES = {
-        ClusterType.AWS: "checkpoints",
-        ClusterType.FAIR: "checkpoint",
-        ClusterType.RSC: "checkpoint/dino",
-    }
-    return Path("/") / CHECKPOINT_DIRNAMES[cluster_type]
-
-
-def get_user_checkpoint_path(
-    cluster_type: Optional[ClusterType] = None,
-) -> Optional[Path]:
-    checkpoint_path = get_checkpoint_path(cluster_type)
-    if checkpoint_path is None:
-        return None
-
-    username = os.environ.get("USER")
-    assert username is not None
-    return checkpoint_path / username
-
-
-def get_slurm_partition(cluster_type: Optional[ClusterType] = None) -> Optional[str]:
-    cluster_type = get_cluster_type(cluster_type)
-    if cluster_type is None:
-        return None
-
-    SLURM_PARTITIONS = {
-        ClusterType.AWS: "learnlab",
-        ClusterType.FAIR: "learnlab",
-        ClusterType.RSC: "learn",
-    }
-    return SLURM_PARTITIONS[cluster_type]
-
-
-def get_slurm_executor_parameters(
-    nodes: int,
-    num_gpus_per_node: int,
-    cluster_type: Optional[ClusterType] = None,
-    **kwargs,
-) -> Dict[str, Any]:
-    # create default parameters
-    params = {
-        "mem_gb": 0,  # Requests all memory on a node, see https://slurm.schedmd.com/sbatch.html
-        "gpus_per_node": num_gpus_per_node,
-        "tasks_per_node": num_gpus_per_node,  # one task per GPU
-        "cpus_per_task": 10,
-        "nodes": nodes,
-        "slurm_partition": get_slurm_partition(cluster_type),
-    }
-    # apply cluster-specific adjustments
-    cluster_type = get_cluster_type(cluster_type)
-    if cluster_type == ClusterType.AWS:
-        params["cpus_per_task"] = 12
-        del params["mem_gb"]
-    elif cluster_type == ClusterType.RSC:
-        params["cpus_per_task"] = 12
-    # set additional parameters / apply overrides
-    params.update(kwargs)
-    return params

mapanything/models/external/dinov2/utils/config.py

@@ -1,74 +0,0 @@
-# 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.
-
-import logging
-import math
-import os
-
-import dinov2.distributed as distributed
-from dinov2.configs import dinov2_default_config
-from dinov2.logging import setup_logging
-from dinov2.utils import utils
-from omegaconf import OmegaConf
-
-logger = logging.getLogger("dinov2")
-
-
-def apply_scaling_rules_to_cfg(cfg):  # to fix
-    if cfg.optim.scaling_rule == "sqrt_wrt_1024":
-        base_lr = cfg.optim.base_lr
-        cfg.optim.lr = base_lr
-        cfg.optim.lr *= math.sqrt(
-            cfg.train.batch_size_per_gpu * distributed.get_global_size() / 1024.0
-        )
-        logger.info(f"sqrt scaling learning rate; base: {base_lr}, new: {cfg.optim.lr}")
-    else:
-        raise NotImplementedError
-    return cfg
-
-
-def write_config(cfg, output_dir, name="config.yaml"):
-    logger.info(OmegaConf.to_yaml(cfg))
-    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_cfg_from_args(args):
-    args.output_dir = os.path.abspath(args.output_dir)
-    args.opts += [f"train.output_dir={args.output_dir}"]
-    default_cfg = OmegaConf.create(dinov2_default_config)
-    cfg = OmegaConf.load(args.config_file)
-    cfg = OmegaConf.merge(default_cfg, cfg, OmegaConf.from_cli(args.opts))
-    return cfg
-
-
-def default_setup(args):
-    distributed.enable(overwrite=True)
-    seed = getattr(args, "seed", 0)
-    rank = distributed.get_global_rank()
-
-    global logger
-    setup_logging(output=args.output_dir, level=logging.INFO)
-    logger = logging.getLogger("dinov2")
-
-    utils.fix_random_seeds(seed + rank)
-    logger.info("git:\n  {}\n".format(utils.get_sha()))
-    logger.info(
-        "\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items()))
-    )
-
-
-def setup(args):
-    """
-    Create configs and perform basic setups.
-    """
-    cfg = get_cfg_from_args(args)
-    os.makedirs(args.output_dir, exist_ok=True)
-    default_setup(args)
-    apply_scaling_rules_to_cfg(cfg)
-    write_config(cfg, args.output_dir)
-    return cfg

mapanything/models/external/dinov2/utils/dtype.py

@@ -1,38 +0,0 @@
-# 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.
-
-
-from typing import Dict, Union
-
-import numpy as np
-import torch
-
-TypeSpec = Union[str, np.dtype, torch.dtype]
-
-
-_NUMPY_TO_TORCH_DTYPE: Dict[np.dtype, torch.dtype] = {
-    np.dtype("bool"): torch.bool,
-    np.dtype("uint8"): torch.uint8,
-    np.dtype("int8"): torch.int8,
-    np.dtype("int16"): torch.int16,
-    np.dtype("int32"): torch.int32,
-    np.dtype("int64"): torch.int64,
-    np.dtype("float16"): torch.float16,
-    np.dtype("float32"): torch.float32,
-    np.dtype("float64"): torch.float64,
-    np.dtype("complex64"): torch.complex64,
-    np.dtype("complex128"): torch.complex128,
-}
-
-
-def as_torch_dtype(dtype: TypeSpec) -> torch.dtype:
-    if isinstance(dtype, torch.dtype):
-        return dtype
-    if isinstance(dtype, str):
-        dtype = np.dtype(dtype)
-    assert isinstance(dtype, np.dtype), (
-        f"Expected an instance of nunpy dtype, got {type(dtype)}"
-    )
-    return _NUMPY_TO_TORCH_DTYPE[dtype]

mapanything/models/external/dinov2/utils/param_groups.py

@@ -1,122 +0,0 @@
-# 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.
-
-import logging
-from collections import defaultdict
-
-logger = logging.getLogger("dinov2")
-
-
-def get_vit_lr_decay_rate(
-    name,
-    lr_decay_rate=1.0,
-    num_layers=12,
-    force_is_backbone=False,
-    chunked_blocks=False,
-):
-    """
-    Calculate lr decay rate for different ViT blocks.
-    Args:
-        name (string): parameter name.
-        lr_decay_rate (float): base lr decay rate.
-        num_layers (int): number of ViT blocks.
-    Returns:
-        lr decay rate for the given parameter.
-    """
-    layer_id = num_layers + 1
-    if name.startswith("backbone") or force_is_backbone:
-        if (
-            ".pos_embed" in name
-            or ".patch_embed" in name
-            or ".mask_token" in name
-            or ".cls_token" in name
-            or ".register_tokens" in name
-        ):
-            layer_id = 0
-        elif force_is_backbone and (
-            "pos_embed" in name
-            or "patch_embed" in name
-            or "mask_token" in name
-            or "cls_token" in name
-            or "register_tokens" in name
-        ):
-            layer_id = 0
-        elif ".blocks." in name and ".residual." not in name:
-            layer_id = int(name[name.find(".blocks.") :].split(".")[2]) + 1
-        elif chunked_blocks and "blocks." in name and "residual." not in name:
-            layer_id = int(name[name.find("blocks.") :].split(".")[2]) + 1
-        elif "blocks." in name and "residual." not in name:
-            layer_id = int(name[name.find("blocks.") :].split(".")[1]) + 1
-
-    return lr_decay_rate ** (num_layers + 1 - layer_id)
-
-
-def get_params_groups_with_decay(model, lr_decay_rate=1.0, patch_embed_lr_mult=1.0):
-    chunked_blocks = False
-    if hasattr(model, "n_blocks"):
-        logger.info("chunked fsdp")
-        n_blocks = model.n_blocks
-        chunked_blocks = model.chunked_blocks
-    elif hasattr(model, "blocks"):
-        logger.info("first code branch")
-        n_blocks = len(model.blocks)
-    elif hasattr(model, "backbone"):
-        logger.info("second code branch")
-        n_blocks = len(model.backbone.blocks)
-    else:
-        logger.info("else code branch")
-        n_blocks = 0
-    all_param_groups = []
-
-    for name, param in model.named_parameters():
-        name = name.replace("_fsdp_wrapped_module.", "")
-        if not param.requires_grad:
-            continue
-        decay_rate = get_vit_lr_decay_rate(
-            name,
-            lr_decay_rate,
-            num_layers=n_blocks,
-            force_is_backbone=n_blocks > 0,
-            chunked_blocks=chunked_blocks,
-        )
-        d = {
-            "params": param,
-            "is_last_layer": False,
-            "lr_multiplier": decay_rate,
-            "wd_multiplier": 1.0,
-            "name": name,
-        }
-
-        if "last_layer" in name:
-            d.update({"is_last_layer": True})
-
-        if name.endswith(".bias") or "norm" in name or "gamma" in name:
-            d.update({"wd_multiplier": 0.0})
-
-        if "patch_embed" in name:
-            d.update({"lr_multiplier": d["lr_multiplier"] * patch_embed_lr_mult})
-
-        all_param_groups.append(d)
-        logger.info(
-            f"""{name}: lr_multiplier: {d["lr_multiplier"]}, wd_multiplier: {d["wd_multiplier"]}"""
-        )
-
-    return all_param_groups
-
-
-def fuse_params_groups(
-    all_params_groups, keys=("lr_multiplier", "wd_multiplier", "is_last_layer")
-):
-    fused_params_groups = defaultdict(lambda: {"params": []})
-    for d in all_params_groups:
-        identifier = ""
-        for k in keys:
-            identifier += k + str(d[k]) + "_"
-
-        for k in keys:
-            fused_params_groups[identifier][k] = d[k]
-        fused_params_groups[identifier]["params"].append(d["params"])
-
-    return fused_params_groups.values()

mapanything/models/external/dinov2/utils/utils.py

@@ -1,105 +0,0 @@
-# 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.
-
-import os
-import random
-import subprocess
-from urllib.parse import urlparse
-
-import numpy as np
-import torch
-from torch import nn
-
-# logger = logging.getLogger("dinov2")
-
-
-def load_pretrained_weights(model, pretrained_weights, checkpoint_key):
-    if urlparse(pretrained_weights).scheme:  # If it looks like an URL
-        state_dict = torch.hub.load_state_dict_from_url(
-            pretrained_weights, map_location="cpu"
-        )
-    else:
-        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()}
-    _ = model.load_state_dict(state_dict, strict=False)
-    # logger.info("Pretrained weights found at {} and loaded with msg: {}".format(pretrained_weights, msg))
-
-
-def fix_random_seeds(seed=31):
-    """
-    Fix random seeds.
-    """
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    np.random.seed(seed)
-    random.seed(seed)
-
-
-def get_sha():
-    cwd = os.path.dirname(os.path.abspath(__file__))
-
-    def _run(command):
-        return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
-
-    sha = "N/A"
-    diff = "clean"
-    branch = "N/A"
-    try:
-        sha = _run(["git", "rev-parse", "HEAD"])
-        subprocess.check_output(["git", "diff"], cwd=cwd)
-        diff = _run(["git", "diff-index", "HEAD"])
-        diff = "has uncommitted changes" if diff else "clean"
-        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
-    except Exception:
-        pass
-    message = f"sha: {sha}, status: {diff}, branch: {branch}"
-    return message
-
-
-class CosineScheduler(object):
-    def __init__(
-        self,
-        base_value,
-        final_value,
-        total_iters,
-        warmup_iters=0,
-        start_warmup_value=0,
-        freeze_iters=0,
-    ):
-        super().__init__()
-        self.final_value = final_value
-        self.total_iters = total_iters
-
-        freeze_schedule = np.zeros((freeze_iters))
-
-        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
-
-        iters = np.arange(total_iters - warmup_iters - freeze_iters)
-        schedule = final_value + 0.5 * (base_value - final_value) * (
-            1 + np.cos(np.pi * iters / len(iters))
-        )
-        self.schedule = np.concatenate((freeze_schedule, warmup_schedule, schedule))
-
-        assert len(self.schedule) == self.total_iters
-
-    def __getitem__(self, it):
-        if it >= self.total_iters:
-            return self.final_value
-        else:
-            return self.schedule[it]
-
-
-def has_batchnorms(model):
-    bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
-    for name, module in model.named_modules():
-        if isinstance(module, bn_types):
-            return True
-    return False

mapanything/models/external/dust3r/__init__.py

@@ -1,217 +0,0 @@
-"""
-Inference wrapper for DUSt3R
-"""
-
-import warnings
-
-import torch
-from dust3r.cloud_opt import global_aligner, GlobalAlignerMode
-from dust3r.image_pairs import make_pairs
-from dust3r.inference import inference
-from dust3r.model import AsymmetricCroCo3DStereo  # noqa
-
-from mapanything.models.external.vggt.utils.rotation import mat_to_quat
-from mapanything.utils.geometry import (
-    convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap,
-    convert_z_depth_to_depth_along_ray,
-    depthmap_to_camera_frame,
-    get_rays_in_camera_frame,
-)
-
-inf = float("inf")
-
-
-def load_model(model_path, device, verbose=True):
-    if verbose:
-        print("Loading model from", model_path)
-    ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
-    args = ckpt["args"].model.replace("ManyAR_PatchEmbed", "PatchEmbedDust3R")
-    if "landscape_only" not in args:
-        args = args[:-1] + ", landscape_only=False)"
-    else:
-        args = args.replace(" ", "").replace(
-            "landscape_only=True", "landscape_only=False"
-        )
-    assert "landscape_only=False" in args
-    if verbose:
-        print(f"Instantiating: {args}")
-    try:
-        net = eval(args)
-    except NameError:
-        net = AsymmetricCroCo3DStereo(
-            enc_depth=24,
-            dec_depth=12,
-            enc_embed_dim=1024,
-            dec_embed_dim=768,
-            enc_num_heads=16,
-            dec_num_heads=12,
-            pos_embed="RoPE100",
-            patch_embed_cls="PatchEmbedDust3R",
-            img_size=(512, 512),
-            head_type="dpt",
-            output_mode="pts3d",
-            depth_mode=("exp", -inf, inf),
-            conf_mode=("exp", 1, inf),
-            landscape_only=False,
-        )
-    s = net.load_state_dict(ckpt["model"], strict=False)
-    if verbose:
-        print(s)
-    return net.to(device)
-
-
-class DUSt3RBAWrapper(torch.nn.Module):
-    def __init__(
-        self,
-        name,
-        ckpt_path,
-        scene_graph="complete",
-        inference_batch_size=32,
-        global_optim_schedule="cosine",
-        global_optim_lr=0.01,
-        global_optim_niter=300,
-        **kwargs,
-    ):
-        super().__init__()
-        self.name = name
-        self.ckpt_path = ckpt_path
-        self.scene_graph = scene_graph
-        self.inference_batch_size = inference_batch_size
-        self.global_optim_schedule = global_optim_schedule
-        self.global_optim_lr = global_optim_lr
-        self.global_optim_niter = global_optim_niter
-
-        # Init the model and load the checkpoint
-        self.model = load_model(self.ckpt_path, device="cpu")
-
-        # Init the global aligner mode
-        self.global_aligner_mode = GlobalAlignerMode.PointCloudOptimizer
-
-    def forward(self, views):
-        """
-        Forward pass wrapper for DUSt3R using the global aligner.
-
-        Assumption:
-        - The batch size of input views is 1.
-
-        Args:
-            views (List[dict]): List of dictionaries containing the input views' images and instance information.
-                                Each dictionary should contain the following keys, where B is the batch size and is 1:
-                                    "img" (tensor): Image tensor of shape (B, C, H, W).
-                                    "data_norm_type" (list): ["dust3r"]
-
-        Returns:
-            List[dict]: A list containing the final outputs for the input views.
-        """
-        # Check the batch size of input views
-        batch_size_per_view, _, height, width = views[0]["img"].shape
-        device = views[0]["img"].device
-        num_views = len(views)
-        assert batch_size_per_view == 1, (
-            f"Batch size of input views should be 1, but got {batch_size_per_view}."
-        )
-
-        # Check the data norm type
-        data_norm_type = views[0]["data_norm_type"][0]
-        assert data_norm_type == "dust3r", (
-            "DUSt3R expects a normalized image with the DUSt3R normalization scheme applied"
-        )
-
-        # Convert the input views to the expected input format
-        images = []
-        for view in views:
-            images.append(
-                dict(
-                    img=view["img"],
-                    idx=len(images),
-                    instance=str(len(images)),
-                )
-            )
-
-        # Make image pairs and run inference pair-wise
-        pairs = make_pairs(
-            images, scene_graph=self.scene_graph, prefilter=None, symmetrize=True
-        )
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore", category=FutureWarning)
-            output = inference(
-                pairs,
-                self.model,
-                device,
-                batch_size=self.inference_batch_size,
-                verbose=False,
-            )
-
-        # Global optimization
-        with torch.enable_grad():
-            scene = global_aligner(
-                output, device=device, mode=self.global_aligner_mode, verbose=False
-            )
-            _ = scene.compute_global_alignment(
-                init="mst",
-                niter=self.global_optim_niter,
-                schedule=self.global_optim_schedule,
-                lr=self.global_optim_lr,
-            )
-
-        # Make sure scene is not None
-        if scene is None:
-            raise RuntimeError("Global optimization failed.")
-
-        # Get the predictions
-        intrinsics = scene.get_intrinsics()
-        c2w_poses = scene.get_im_poses()
-        depths = scene.get_depthmaps()
-
-        # Convert the output to the MapAnything format
-        with torch.autocast("cuda", enabled=False):
-            res = []
-            for view_idx in range(num_views):
-                # Get the current view predictions
-                curr_view_intrinsic = intrinsics[view_idx].unsqueeze(0)
-                curr_view_pose = c2w_poses[view_idx].unsqueeze(0)
-                curr_view_depth_z = depths[view_idx].unsqueeze(0)
-
-                # Convert the pose to quaternions and translation
-                curr_view_cam_translations = curr_view_pose[..., :3, 3]
-                curr_view_cam_quats = mat_to_quat(curr_view_pose[..., :3, :3])
-
-                # Get the camera frame pointmaps
-                curr_view_pts3d_cam, _ = depthmap_to_camera_frame(
-                    curr_view_depth_z, curr_view_intrinsic
-                )
-
-                # Convert the z depth to depth along ray
-                curr_view_depth_along_ray = convert_z_depth_to_depth_along_ray(
-                    curr_view_depth_z, curr_view_intrinsic
-                )
-                curr_view_depth_along_ray = curr_view_depth_along_ray.unsqueeze(-1)
-
-                # Get the ray directions on the unit sphere in the camera frame
-                _, curr_view_ray_dirs = get_rays_in_camera_frame(
-                    curr_view_intrinsic, height, width, normalize_to_unit_sphere=True
-                )
-
-                # Get the pointmaps
-                curr_view_pts3d = (
-                    convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap(
-                        curr_view_ray_dirs,
-                        curr_view_depth_along_ray,
-                        curr_view_cam_translations,
-                        curr_view_cam_quats,
-                    )
-                )
-
-                # Append the outputs to the result list
-                res.append(
-                    {
-                        "pts3d": curr_view_pts3d,
-                        "pts3d_cam": curr_view_pts3d_cam,
-                        "ray_directions": curr_view_ray_dirs,
-                        "depth_along_ray": curr_view_depth_along_ray,
-                        "cam_trans": curr_view_cam_translations,
-                        "cam_quats": curr_view_cam_quats,
-                    }
-                )
-
-        return res

mapanything/models/external/mast3r/__init__.py

@@ -1,191 +0,0 @@
-"""
-Inference wrapper for MASt3R + Sparse GA
-"""
-
-import os
-import tempfile
-import warnings
-
-import torch
-from dust3r.image_pairs import make_pairs
-from mast3r.cloud_opt.sparse_ga import sparse_global_alignment
-from mast3r.model import load_model
-
-from mapanything.models.external.vggt.utils.rotation import mat_to_quat
-from mapanything.utils.geometry import (
-    convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap,
-    convert_z_depth_to_depth_along_ray,
-    depthmap_to_camera_frame,
-    get_rays_in_camera_frame,
-)
-
-
-class MASt3RSGAWrapper(torch.nn.Module):
-    def __init__(
-        self,
-        name,
-        ckpt_path,
-        cache_dir,
-        scene_graph="complete",
-        sparse_ga_lr1=0.07,
-        sparse_ga_niter1=300,
-        sparse_ga_lr2=0.01,
-        sparse_ga_niter2=300,
-        sparse_ga_optim_level="refine+depth",
-        sparse_ga_shared_intrinsics=False,
-        sparse_ga_matching_conf_thr=5.0,
-        **kwargs,
-    ):
-        super().__init__()
-        self.name = name
-        self.ckpt_path = ckpt_path
-        self.cache_dir = cache_dir
-        self.scene_graph = scene_graph
-        self.sparse_ga_lr1 = sparse_ga_lr1
-        self.sparse_ga_niter1 = sparse_ga_niter1
-        self.sparse_ga_lr2 = sparse_ga_lr2
-        self.sparse_ga_niter2 = sparse_ga_niter2
-        self.sparse_ga_optim_level = sparse_ga_optim_level
-        self.sparse_ga_shared_intrinsics = sparse_ga_shared_intrinsics
-        self.sparse_ga_matching_conf_thr = sparse_ga_matching_conf_thr
-
-        # Init the model and load the checkpoint
-        self.model = load_model(self.ckpt_path, device="cpu")
-
-    def forward(self, views):
-        """
-        Forward pass wrapper for MASt3R using the sparse global aligner.
-
-        Assumption:
-        - The batch size of input views is 1.
-
-        Args:
-            views (List[dict]): List of dictionaries containing the input views' images and instance information.
-                                Each dictionary should contain the following keys, where B is the batch size and is 1:
-                                    "img" (tensor): Image tensor of shape (B, C, H, W).
-                                    "data_norm_type" (list): ["dust3r"]
-                                    "label" (list): ["scene_name"]
-                                    "instance" (list): ["image_name"]
-
-        Returns:
-            List[dict]: A list containing the final outputs for the input views.
-        """
-        # Check the batch size of input views
-        batch_size_per_view, _, height, width = views[0]["img"].shape
-        device = views[0]["img"].device
-        num_views = len(views)
-        assert batch_size_per_view == 1, (
-            f"Batch size of input views should be 1, but got {batch_size_per_view}."
-        )
-
-        # Check the data norm type
-        data_norm_type = views[0]["data_norm_type"][0]
-        assert data_norm_type == "dust3r", (
-            "MASt3R expects a normalized image with the DUSt3R normalization scheme applied"
-        )
-
-        # Convert the input views to the expected input format
-        images = []
-        image_paths = []
-        for view in views:
-            images.append(
-                dict(
-                    img=view["img"].cpu(),
-                    idx=len(images),
-                    instance=str(len(images)),
-                    true_shape=torch.tensor(view["img"].shape[-2:])[None]
-                    .repeat(batch_size_per_view, 1)
-                    .numpy(),
-                )
-            )
-            view_name = os.path.join(view["label"][0], view["instance"][0])
-            image_paths.append(view_name)
-
-        # Make image pairs and run inference
-        # Sparse GA (forward mast3r -> matching -> 3D optim -> 2D refinement -> triangulation)
-        pairs = make_pairs(
-            images, scene_graph=self.scene_graph, prefilter=None, symmetrize=True
-        )
-        with torch.enable_grad():
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore", category=FutureWarning)
-                tempfile.mkdtemp(dir=self.cache_dir)
-                scene = sparse_global_alignment(
-                    image_paths,
-                    pairs,
-                    self.cache_dir,
-                    self.model,
-                    lr1=self.sparse_ga_lr1,
-                    niter1=self.sparse_ga_niter1,
-                    lr2=self.sparse_ga_lr2,
-                    niter2=self.sparse_ga_niter2,
-                    device=device,
-                    opt_depth="depth" in self.sparse_ga_optim_level,
-                    shared_intrinsics=self.sparse_ga_shared_intrinsics,
-                    matching_conf_thr=self.sparse_ga_matching_conf_thr,
-                    verbose=False,
-                )
-
-        # Make sure scene is not None
-        if scene is None:
-            raise RuntimeError("Global optimization failed.")
-
-        # Get the predictions
-        intrinsics = scene.intrinsics
-        c2w_poses = scene.get_im_poses()
-        _, depths, _ = scene.get_dense_pts3d()
-
-        # Convert the output to the MapAnything format
-        with torch.autocast("cuda", enabled=False):
-            res = []
-            for view_idx in range(num_views):
-                # Get the current view predictions
-                curr_view_intrinsic = intrinsics[view_idx].unsqueeze(0)
-                curr_view_pose = c2w_poses[view_idx].unsqueeze(0)
-                curr_view_depth_z = (
-                    depths[view_idx].reshape((height, width)).unsqueeze(0)
-                )
-
-                # Convert the pose to quaternions and translation
-                curr_view_cam_translations = curr_view_pose[..., :3, 3]
-                curr_view_cam_quats = mat_to_quat(curr_view_pose[..., :3, :3])
-
-                # Get the camera frame pointmaps
-                curr_view_pts3d_cam, _ = depthmap_to_camera_frame(
-                    curr_view_depth_z, curr_view_intrinsic
-                )
-
-                # Convert the z depth to depth along ray
-                curr_view_depth_along_ray = convert_z_depth_to_depth_along_ray(
-                    curr_view_depth_z, curr_view_intrinsic
-                )
-                curr_view_depth_along_ray = curr_view_depth_along_ray.unsqueeze(-1)
-
-                # Get the ray directions on the unit sphere in the camera frame
-                _, curr_view_ray_dirs = get_rays_in_camera_frame(
-                    curr_view_intrinsic, height, width, normalize_to_unit_sphere=True
-                )
-
-                # Get the pointmaps
-                curr_view_pts3d = (
-                    convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap(
-                        curr_view_ray_dirs,
-                        curr_view_depth_along_ray,
-                        curr_view_cam_translations,
-                        curr_view_cam_quats,
-                    )
-                )
-
-                # Append the outputs to the result list
-                res.append(
-                    {
-                        "pts3d": curr_view_pts3d,
-                        "pts3d_cam": curr_view_pts3d_cam,
-                        "ray_directions": curr_view_ray_dirs,
-                        "depth_along_ray": curr_view_depth_along_ray,
-                        "cam_trans": curr_view_cam_translations,
-                        "cam_quats": curr_view_cam_quats,
-                    }
-                )
-
-        return res

mapanything/models/external/moge/__init__.py

@@ -1,114 +0,0 @@
-"""
-Inference wrapper for MoGe
-"""
-
-import torch
-
-from mapanything.models.external.moge.models.v1 import MoGeModel as MoGeModelV1
-from mapanything.models.external.moge.models.v2 import MoGeModel as MoGeModelV2
-
-
-class MoGeWrapper(torch.nn.Module):
-    def __init__(
-        self,
-        name,
-        model_string="Ruicheng/moge-2-vitl",
-        torch_hub_force_reload=False,
-        load_custom_ckpt=False,
-        custom_ckpt_path=None,
-    ):
-        super().__init__()
-        self.name = name
-        self.model_string = model_string
-        self.torch_hub_force_reload = torch_hub_force_reload
-        self.load_custom_ckpt = load_custom_ckpt
-        self.custom_ckpt_path = custom_ckpt_path
-
-        # Mapping of MoGe model version to checkpoint strings
-        self.moge_model_map = {
-            "v1": ["Ruicheng/moge-vitl"],
-            "v2": [
-                "Ruicheng/moge-2-vits-normal",
-                "Ruicheng/moge-2-vitb-normal",
-                "Ruicheng/moge-2-vitl-normal",
-                "Ruicheng/moge-2-vitl",
-            ],
-        }
-
-        # Initialize the model
-        if self.model_string in self.moge_model_map["v1"]:
-            self.model = MoGeModelV1.from_pretrained(self.model_string)
-        elif self.model_string in self.moge_model_map["v2"]:
-            self.model = MoGeModelV2.from_pretrained(self.model_string)
-        else:
-            raise ValueError(
-                f"Invalid model string: {self.model_string}. Valid strings are: {self.moge_model_map}"
-            )
-
-        # Load custom checkpoint if requested
-        if self.load_custom_ckpt:
-            print(f"Loading checkpoint from {self.custom_ckpt_path} ...")
-            assert self.custom_ckpt_path is not None, (
-                "custom_ckpt_path must be provided if load_custom_ckpt is set to True"
-            )
-            custom_ckpt = torch.load(self.custom_ckpt_path, weights_only=False)
-            print(self.model.load_state_dict(custom_ckpt, strict=True))
-            del custom_ckpt  # in case it occupies memory
-
-    def forward(self, views):
-        """
-        Forward pass wrapper for MoGe-2.
-        The predicted MoGe-2 mask is not applied to the outputs.
-        The number of tokens for inference is determined by the image shape.
-
-        Assumption:
-        - The number of input views is 1.
-
-        Args:
-            views (List[dict]): List of dictionaries containing the input views' images and instance information.
-                                Length of the list should be 1.
-                                Each dictionary should contain the following keys:
-                                    "img" (tensor): Image tensor of shape (B, C, H, W).
-                                    "data_norm_type" (list): ["identity"]
-
-        Returns:
-            List[dict]: A list containing the final outputs for the single view. Length of the list will be 1.
-        """
-        # Check that the number of input views is 1
-        assert len(views) == 1, "MoGe only supports 1 input view."
-
-        # Get input shape of the images, number of tokens for inference, and batch size per view
-        _, _, height, width = views[0]["img"].shape
-        num_tokens = int(height // 14) * int(width // 14)
-
-        # Check the data norm type
-        # MoGe expects a normalized image but without the DINOv2 mean and std applied ("identity")
-        data_norm_type = views[0]["data_norm_type"][0]
-        assert data_norm_type == "identity", (
-            "MoGe expects a normalized image but without the DINOv2 mean and std applied"
-        )
-
-        # Run MoGe inference
-        # Output dict contains: "points", "depth", "mask", "intrinsics", "normal" (based on model config)
-        model_outputs = self.model.infer(
-            image=views[0]["img"], num_tokens=num_tokens, apply_mask=False
-        )
-
-        # Get the ray directions and depth along ray
-        with torch.autocast("cuda", enabled=False):
-            depth_along_ray = torch.norm(model_outputs["points"], dim=-1, keepdim=True)
-            ray_directions = model_outputs["points"] / depth_along_ray
-
-        # Convert the output to MapAnything format
-        result_dict = {
-            "pts3d": model_outputs["points"],
-            "pts3d_cam": model_outputs["points"],
-            "depth_z": model_outputs["depth"].unsqueeze(-1),
-            "intrinsics": model_outputs["intrinsics"],
-            "non_ambiguous_mask": model_outputs["mask"],
-            "ray_directions": ray_directions,
-            "depth_along_ray": depth_along_ray,
-        }
-        res = [result_dict]
-
-        return res

mapanything/models/external/moge/models/modules.py

@@ -1,467 +0,0 @@
-import functools
-import importlib
-import itertools
-from typing import List, Literal, Optional, Sequence, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from mapanything.models.external.dinov2.models.vision_transformer import (
-    DinoVisionTransformer,
-)
-from mapanything.models.external.moge.models.utils import (
-    wrap_dinov2_attention_with_sdpa,
-    wrap_module_with_gradient_checkpointing,
-)
-
-
-class ResidualConvBlock(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int = None,
-        hidden_channels: int = None,
-        kernel_size: int = 3,
-        padding_mode: str = "replicate",
-        activation: Literal["relu", "leaky_relu", "silu", "elu"] = "relu",
-        in_norm: Literal[
-            "group_norm", "layer_norm", "instance_norm", "none"
-        ] = "layer_norm",
-        hidden_norm: Literal[
-            "group_norm", "layer_norm", "instance_norm"
-        ] = "group_norm",
-    ):
-        super(ResidualConvBlock, self).__init__()
-        if out_channels is None:
-            out_channels = in_channels
-        if hidden_channels is None:
-            hidden_channels = in_channels
-
-        if activation == "relu":
-            activation_cls = nn.ReLU
-        elif activation == "leaky_relu":
-            activation_cls = functools.partial(nn.LeakyReLU, negative_slope=0.2)
-        elif activation == "silu":
-            activation_cls = nn.SiLU
-        elif activation == "elu":
-            activation_cls = nn.ELU
-        else:
-            raise ValueError(f"Unsupported activation function: {activation}")
-
-        self.layers = nn.Sequential(
-            nn.GroupNorm(in_channels // 32, in_channels)
-            if in_norm == "group_norm"
-            else nn.GroupNorm(1, in_channels)
-            if in_norm == "layer_norm"
-            else nn.InstanceNorm2d(in_channels)
-            if in_norm == "instance_norm"
-            else nn.Identity(),
-            activation_cls(),
-            nn.Conv2d(
-                in_channels,
-                hidden_channels,
-                kernel_size=kernel_size,
-                padding=kernel_size // 2,
-                padding_mode=padding_mode,
-            ),
-            nn.GroupNorm(hidden_channels // 32, hidden_channels)
-            if hidden_norm == "group_norm"
-            else nn.GroupNorm(1, hidden_channels)
-            if hidden_norm == "layer_norm"
-            else nn.InstanceNorm2d(hidden_channels)
-            if hidden_norm == "instance_norm"
-            else nn.Identity(),
-            activation_cls(),
-            nn.Conv2d(
-                hidden_channels,
-                out_channels,
-                kernel_size=kernel_size,
-                padding=kernel_size // 2,
-                padding_mode=padding_mode,
-            ),
-        )
-
-        self.skip_connection = (
-            nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)
-            if in_channels != out_channels
-            else nn.Identity()
-        )
-
-    def forward(self, x):
-        skip = self.skip_connection(x)
-        x = self.layers(x)
-        x = x + skip
-        return x
-
-
-class DINOv2Encoder(nn.Module):
-    "Wrapped DINOv2 encoder supporting gradient checkpointing. Input is RGB image in range [0, 1]."
-
-    backbone: DinoVisionTransformer
-    image_mean: torch.Tensor
-    image_std: torch.Tensor
-    dim_features: int
-
-    def __init__(
-        self,
-        backbone: str,
-        intermediate_layers: Union[int, List[int]],
-        dim_out: int,
-        **deprecated_kwargs,
-    ):
-        super(DINOv2Encoder, self).__init__()
-
-        self.intermediate_layers = intermediate_layers
-
-        # Load the backbone
-        self.hub_loader = getattr(
-            importlib.import_module(
-                "mapanything.models.external.dinov2.hub.backbones", __package__
-            ),
-            backbone,
-        )
-        self.backbone_name = backbone
-        self.backbone = self.hub_loader(pretrained=False)
-
-        self.dim_features = self.backbone.blocks[0].attn.qkv.in_features
-        self.num_features = (
-            intermediate_layers
-            if isinstance(intermediate_layers, int)
-            else len(intermediate_layers)
-        )
-
-        self.output_projections = nn.ModuleList(
-            [
-                nn.Conv2d(
-                    in_channels=self.dim_features,
-                    out_channels=dim_out,
-                    kernel_size=1,
-                    stride=1,
-                    padding=0,
-                )
-                for _ in range(self.num_features)
-            ]
-        )
-
-        self.register_buffer(
-            "image_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
-        )
-        self.register_buffer(
-            "image_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
-        )
-
-    @property
-    def onnx_compatible_mode(self):
-        return getattr(self, "_onnx_compatible_mode", False)
-
-    @onnx_compatible_mode.setter
-    def onnx_compatible_mode(self, value: bool):
-        self._onnx_compatible_mode = value
-        self.backbone.onnx_compatible_mode = value
-
-    def init_weights(self):
-        pretrained_backbone_state_dict = self.hub_loader(pretrained=True).state_dict()
-        self.backbone.load_state_dict(pretrained_backbone_state_dict)
-
-    def enable_gradient_checkpointing(self):
-        for i in range(len(self.backbone.blocks)):
-            wrap_module_with_gradient_checkpointing(self.backbone.blocks[i])
-
-    def enable_pytorch_native_sdpa(self):
-        for i in range(len(self.backbone.blocks)):
-            wrap_dinov2_attention_with_sdpa(self.backbone.blocks[i].attn)
-
-    def forward(
-        self,
-        image: torch.Tensor,
-        token_rows: Union[int, torch.LongTensor],
-        token_cols: Union[int, torch.LongTensor],
-        return_class_token: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        image_14 = F.interpolate(
-            image,
-            (token_rows * 14, token_cols * 14),
-            mode="bilinear",
-            align_corners=False,
-            antialias=not self.onnx_compatible_mode,
-        )
-        image_14 = (image_14 - self.image_mean) / self.image_std
-
-        # Get intermediate layers from the backbone
-        features = self.backbone.get_intermediate_layers(
-            image_14, n=self.intermediate_layers, return_class_token=True
-        )
-
-        # Project features to the desired dimensionality
-        x = torch.stack(
-            [
-                proj(
-                    feat.permute(0, 2, 1)
-                    .unflatten(2, (token_rows, token_cols))
-                    .contiguous()
-                )
-                for proj, (feat, clstoken) in zip(self.output_projections, features)
-            ],
-            dim=1,
-        ).sum(dim=1)
-
-        if return_class_token:
-            return x, features[-1][1]
-        else:
-            return x
-
-
-class Resampler(nn.Sequential):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        type_: Literal[
-            "pixel_shuffle",
-            "nearest",
-            "bilinear",
-            "conv_transpose",
-            "pixel_unshuffle",
-            "avg_pool",
-            "max_pool",
-        ],
-        scale_factor: int = 2,
-    ):
-        if type_ == "pixel_shuffle":
-            nn.Sequential.__init__(
-                self,
-                nn.Conv2d(
-                    in_channels,
-                    out_channels * (scale_factor**2),
-                    kernel_size=3,
-                    stride=1,
-                    padding=1,
-                    padding_mode="replicate",
-                ),
-                nn.PixelShuffle(scale_factor),
-                nn.Conv2d(
-                    out_channels,
-                    out_channels,
-                    kernel_size=3,
-                    stride=1,
-                    padding=1,
-                    padding_mode="replicate",
-                ),
-            )
-            for i in range(1, scale_factor**2):
-                self[0].weight.data[i :: scale_factor**2] = self[0].weight.data[
-                    0 :: scale_factor**2
-                ]
-                self[0].bias.data[i :: scale_factor**2] = self[0].bias.data[
-                    0 :: scale_factor**2
-                ]
-        elif type_ in ["nearest", "bilinear"]:
-            nn.Sequential.__init__(
-                self,
-                nn.Upsample(
-                    scale_factor=scale_factor,
-                    mode=type_,
-                    align_corners=False if type_ == "bilinear" else None,
-                ),
-                nn.Conv2d(
-                    in_channels,
-                    out_channels,
-                    kernel_size=3,
-                    stride=1,
-                    padding=1,
-                    padding_mode="replicate",
-                ),
-            )
-        elif type_ == "conv_transpose":
-            nn.Sequential.__init__(
-                self,
-                nn.ConvTranspose2d(
-                    in_channels,
-                    out_channels,
-                    kernel_size=scale_factor,
-                    stride=scale_factor,
-                ),
-                nn.Conv2d(
-                    out_channels,
-                    out_channels,
-                    kernel_size=3,
-                    stride=1,
-                    padding=1,
-                    padding_mode="replicate",
-                ),
-            )
-            self[0].weight.data[:] = self[0].weight.data[:, :, :1, :1]
-        elif type_ == "pixel_unshuffle":
-            nn.Sequential.__init__(
-                self,
-                nn.PixelUnshuffle(scale_factor),
-                nn.Conv2d(
-                    in_channels * (scale_factor**2),
-                    out_channels,
-                    kernel_size=3,
-                    stride=1,
-                    padding=1,
-                    padding_mode="replicate",
-                ),
-            )
-        elif type_ == "avg_pool":
-            nn.Sequential.__init__(
-                self,
-                nn.Conv2d(
-                    in_channels,
-                    out_channels,
-                    kernel_size=3,
-                    stride=1,
-                    padding=1,
-                    padding_mode="replicate",
-                ),
-                nn.AvgPool2d(kernel_size=scale_factor, stride=scale_factor),
-            )
-        elif type_ == "max_pool":
-            nn.Sequential.__init__(
-                self,
-                nn.Conv2d(
-                    in_channels,
-                    out_channels,
-                    kernel_size=3,
-                    stride=1,
-                    padding=1,
-                    padding_mode="replicate",
-                ),
-                nn.MaxPool2d(kernel_size=scale_factor, stride=scale_factor),
-            )
-        else:
-            raise ValueError(f"Unsupported resampler type: {type_}")
-
-
-class MLP(nn.Sequential):
-    def __init__(self, dims: Sequence[int]):
-        nn.Sequential.__init__(
-            self,
-            *itertools.chain(
-                *[
-                    (nn.Linear(dim_in, dim_out), nn.ReLU(inplace=True))
-                    for dim_in, dim_out in zip(dims[:-2], dims[1:-1])
-                ]
-            ),
-            nn.Linear(dims[-2], dims[-1]),
-        )
-
-
-class ConvStack(nn.Module):
-    def __init__(
-        self,
-        dim_in: List[Optional[int]],
-        dim_res_blocks: List[int],
-        dim_out: List[Optional[int]],
-        resamplers: Union[
-            Literal[
-                "pixel_shuffle",
-                "nearest",
-                "bilinear",
-                "conv_transpose",
-                "pixel_unshuffle",
-                "avg_pool",
-                "max_pool",
-            ],
-            List,
-        ],
-        dim_times_res_block_hidden: int = 1,
-        num_res_blocks: int = 1,
-        res_block_in_norm: Literal[
-            "layer_norm", "group_norm", "instance_norm", "none"
-        ] = "layer_norm",
-        res_block_hidden_norm: Literal[
-            "layer_norm", "group_norm", "instance_norm", "none"
-        ] = "group_norm",
-        activation: Literal["relu", "leaky_relu", "silu", "elu"] = "relu",
-    ):
-        super().__init__()
-        self.input_blocks = nn.ModuleList(
-            [
-                nn.Conv2d(dim_in_, dim_res_block_, kernel_size=1, stride=1, padding=0)
-                if dim_in_ is not None
-                else nn.Identity()
-                for dim_in_, dim_res_block_ in zip(
-                    dim_in
-                    if isinstance(dim_in, Sequence)
-                    else itertools.repeat(dim_in),
-                    dim_res_blocks,
-                )
-            ]
-        )
-        self.resamplers = nn.ModuleList(
-            [
-                Resampler(dim_prev, dim_succ, scale_factor=2, type_=resampler)
-                for i, (dim_prev, dim_succ, resampler) in enumerate(
-                    zip(
-                        dim_res_blocks[:-1],
-                        dim_res_blocks[1:],
-                        resamplers
-                        if isinstance(resamplers, Sequence)
-                        else itertools.repeat(resamplers),
-                    )
-                )
-            ]
-        )
-        self.res_blocks = nn.ModuleList(
-            [
-                nn.Sequential(
-                    *(
-                        ResidualConvBlock(
-                            dim_res_block_,
-                            dim_res_block_,
-                            dim_times_res_block_hidden * dim_res_block_,
-                            activation=activation,
-                            in_norm=res_block_in_norm,
-                            hidden_norm=res_block_hidden_norm,
-                        )
-                        for _ in range(
-                            num_res_blocks[i]
-                            if isinstance(num_res_blocks, list)
-                            else num_res_blocks
-                        )
-                    )
-                )
-                for i, dim_res_block_ in enumerate(dim_res_blocks)
-            ]
-        )
-        self.output_blocks = nn.ModuleList(
-            [
-                nn.Conv2d(dim_res_block_, dim_out_, kernel_size=1, stride=1, padding=0)
-                if dim_out_ is not None
-                else nn.Identity()
-                for dim_out_, dim_res_block_ in zip(
-                    dim_out
-                    if isinstance(dim_out, Sequence)
-                    else itertools.repeat(dim_out),
-                    dim_res_blocks,
-                )
-            ]
-        )
-
-    def enable_gradient_checkpointing(self):
-        for i in range(len(self.resamplers)):
-            self.resamplers[i] = wrap_module_with_gradient_checkpointing(
-                self.resamplers[i]
-            )
-        for i in range(len(self.res_blocks)):
-            for j in range(len(self.res_blocks[i])):
-                self.res_blocks[i][j] = wrap_module_with_gradient_checkpointing(
-                    self.res_blocks[i][j]
-                )
-
-    def forward(self, in_features: List[torch.Tensor]):
-        out_features = []
-        for i in range(len(self.res_blocks)):
-            feature = self.input_blocks[i](in_features[i])
-            if i == 0:
-                x = feature
-            elif feature is not None:
-                x = x + feature
-            x = self.res_blocks[i](x)
-            out_features.append(self.output_blocks[i](x))
-            if i < len(self.res_blocks) - 1:
-                x = self.resamplers[i](x)
-        return out_features

mapanything/models/external/moge/models/utils.py

@@ -1,477 +0,0 @@
-import inspect
-from functools import partial, wraps
-from numbers import Number
-from typing import Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-def wrap_module_with_gradient_checkpointing(module: nn.Module):
-    from torch.utils.checkpoint import checkpoint
-
-    class _CheckpointingWrapper(module.__class__):
-        _restore_cls = module.__class__
-
-        def forward(self, *args, **kwargs):
-            return checkpoint(super().forward, *args, use_reentrant=False, **kwargs)
-
-    module.__class__ = _CheckpointingWrapper
-    return module
-
-
-def unwrap_module_with_gradient_checkpointing(module: nn.Module):
-    module.__class__ = module.__class__._restore_cls
-
-
-def wrap_dinov2_attention_with_sdpa(module: nn.Module):
-    assert torch.__version__ >= "2.0", "SDPA requires PyTorch 2.0 or later"
-
-    class _AttentionWrapper(module.__class__):
-        def forward(self, x: torch.Tensor, attn_bias=None) -> torch.Tensor:
-            B, N, C = x.shape
-            qkv = (
-                self.qkv(x)
-                .reshape(B, N, 3, self.num_heads, C // self.num_heads)
-                .permute(2, 0, 3, 1, 4)
-            )  # (3, B, H, N, C // H)
-
-            q, k, v = torch.unbind(qkv, 0)  # (B, H, N, C // H)
-
-            x = F.scaled_dot_product_attention(q, k, v, attn_bias)
-            x = x.permute(0, 2, 1, 3).reshape(B, N, C)
-
-            x = self.proj(x)
-            x = self.proj_drop(x)
-            return x
-
-    module.__class__ = _AttentionWrapper
-    return module
-
-
-def sync_ddp_hook(
-    state, bucket: torch.distributed.GradBucket
-) -> torch.futures.Future[torch.Tensor]:
-    group_to_use = torch.distributed.group.WORLD
-    world_size = group_to_use.size()
-    grad = bucket.buffer()
-    grad.div_(world_size)
-    torch.distributed.all_reduce(grad, group=group_to_use)
-    fut = torch.futures.Future()
-    fut.set_result(grad)
-    return fut
-
-
-def normalized_view_plane_uv(
-    width: int,
-    height: int,
-    aspect_ratio: float = None,
-    dtype: torch.dtype = None,
-    device: torch.device = None,
-) -> torch.Tensor:
-    "UV with left-top corner as (-width / diagonal, -height / diagonal) and right-bottom corner as (width / diagonal, height / diagonal)"
-    if aspect_ratio is None:
-        aspect_ratio = width / height
-
-    span_x = aspect_ratio / (1 + aspect_ratio**2) ** 0.5
-    span_y = 1 / (1 + aspect_ratio**2) ** 0.5
-
-    u = torch.linspace(
-        -span_x * (width - 1) / width,
-        span_x * (width - 1) / width,
-        width,
-        dtype=dtype,
-        device=device,
-    )
-    v = torch.linspace(
-        -span_y * (height - 1) / height,
-        span_y * (height - 1) / height,
-        height,
-        dtype=dtype,
-        device=device,
-    )
-    u, v = torch.meshgrid(u, v, indexing="xy")
-    uv = torch.stack([u, v], dim=-1)
-    return uv
-
-
-def solve_optimal_focal_shift(uv: np.ndarray, xyz: np.ndarray):
-    "Solve `min |focal * xy / (z + shift) - uv|` with respect to shift and focal"
-    from scipy.optimize import least_squares
-
-    uv, xy, z = uv.reshape(-1, 2), xyz[..., :2].reshape(-1, 2), xyz[..., 2].reshape(-1)
-
-    def fn(uv: np.ndarray, xy: np.ndarray, z: np.ndarray, shift: np.ndarray):
-        xy_proj = xy / (z + shift)[:, None]
-        f = (xy_proj * uv).sum() / np.square(xy_proj).sum()
-        err = (f * xy_proj - uv).ravel()
-        return err
-
-    solution = least_squares(partial(fn, uv, xy, z), x0=0, ftol=1e-3, method="lm")
-    optim_shift = solution["x"].squeeze().astype(np.float32)
-
-    xy_proj = xy / (z + optim_shift)[:, None]
-    optim_focal = (xy_proj * uv).sum() / np.square(xy_proj).sum()
-
-    return optim_shift, optim_focal
-
-
-def solve_optimal_shift(uv: np.ndarray, xyz: np.ndarray, focal: float):
-    "Solve `min |focal * xy / (z + shift) - uv|` with respect to shift"
-    from scipy.optimize import least_squares
-
-    uv, xy, z = uv.reshape(-1, 2), xyz[..., :2].reshape(-1, 2), xyz[..., 2].reshape(-1)
-
-    def fn(uv: np.ndarray, xy: np.ndarray, z: np.ndarray, shift: np.ndarray):
-        xy_proj = xy / (z + shift)[:, None]
-        err = (focal * xy_proj - uv).ravel()
-        return err
-
-    solution = least_squares(partial(fn, uv, xy, z), x0=0, ftol=1e-3, method="lm")
-    optim_shift = solution["x"].squeeze().astype(np.float32)
-
-    return optim_shift
-
-
-def recover_focal_shift(
-    points: torch.Tensor,
-    mask: torch.Tensor = None,
-    focal: torch.Tensor = None,
-    downsample_size: Tuple[int, int] = (64, 64),
-):
-    """
-    Recover the depth map and FoV from a point map with unknown z shift and focal.
-
-    Note that it assumes:
-    - the optical center is at the center of the map
-    - the map is undistorted
-    - the map is isometric in the x and y directions
-
-    ### Parameters:
-    - `points: torch.Tensor` of shape (..., H, W, 3)
-    - `downsample_size: Tuple[int, int]` in (height, width), the size of the downsampled map. Downsampling produces approximate solution and is efficient for large maps.
-
-    ### Returns:
-    - `focal`: torch.Tensor of shape (...) the estimated focal length, relative to the half diagonal of the map
-    - `shift`: torch.Tensor of shape (...) Z-axis shift to translate the point map to camera space
-    """
-    shape = points.shape
-    height, width = points.shape[-3], points.shape[-2]
-
-    points = points.reshape(-1, *shape[-3:])
-    mask = None if mask is None else mask.reshape(-1, *shape[-3:-1])
-    focal = focal.reshape(-1) if focal is not None else None
-    uv = normalized_view_plane_uv(
-        width, height, dtype=points.dtype, device=points.device
-    )  # (H, W, 2)
-
-    points_lr = F.interpolate(
-        points.permute(0, 3, 1, 2), downsample_size, mode="nearest"
-    ).permute(0, 2, 3, 1)
-    uv_lr = (
-        F.interpolate(
-            uv.unsqueeze(0).permute(0, 3, 1, 2), downsample_size, mode="nearest"
-        )
-        .squeeze(0)
-        .permute(1, 2, 0)
-    )
-    mask_lr = (
-        None
-        if mask is None
-        else F.interpolate(
-            mask.to(torch.float32).unsqueeze(1), downsample_size, mode="nearest"
-        ).squeeze(1)
-        > 0
-    )
-
-    uv_lr_np = uv_lr.cpu().numpy()
-    points_lr_np = points_lr.detach().cpu().numpy()
-    focal_np = focal.cpu().numpy() if focal is not None else None
-    mask_lr_np = None if mask is None else mask_lr.cpu().numpy()
-    optim_shift, optim_focal = [], []
-    for i in range(points.shape[0]):
-        points_lr_i_np = (
-            points_lr_np[i] if mask is None else points_lr_np[i][mask_lr_np[i]]
-        )
-        uv_lr_i_np = uv_lr_np if mask is None else uv_lr_np[mask_lr_np[i]]
-        if uv_lr_i_np.shape[0] < 2:
-            optim_focal.append(1)
-            optim_shift.append(0)
-            continue
-        if focal is None:
-            optim_shift_i, optim_focal_i = solve_optimal_focal_shift(
-                uv_lr_i_np, points_lr_i_np
-            )
-            optim_focal.append(float(optim_focal_i))
-        else:
-            optim_shift_i = solve_optimal_shift(uv_lr_i_np, points_lr_i_np, focal_np[i])
-        optim_shift.append(float(optim_shift_i))
-    optim_shift = torch.tensor(
-        optim_shift, device=points.device, dtype=points.dtype
-    ).reshape(shape[:-3])
-
-    if focal is None:
-        optim_focal = torch.tensor(
-            optim_focal, device=points.device, dtype=points.dtype
-        ).reshape(shape[:-3])
-    else:
-        optim_focal = focal.reshape(shape[:-3])
-
-    return optim_focal, optim_shift
-
-
-def suppress_traceback(fn):
-    @wraps(fn)
-    def wrapper(*args, **kwargs):
-        try:
-            return fn(*args, **kwargs)
-        except Exception as e:
-            e.__traceback__ = e.__traceback__.tb_next.tb_next
-            raise
-
-    return wrapper
-
-
-def get_device(args, kwargs):
-    device = None
-    for arg in list(args) + list(kwargs.values()):
-        if isinstance(arg, torch.Tensor):
-            if device is None:
-                device = arg.device
-            elif device != arg.device:
-                raise ValueError("All tensors must be on the same device.")
-    return device
-
-
-def get_args_order(func, args, kwargs):
-    """
-    Get the order of the arguments of a function.
-    """
-    names = inspect.getfullargspec(func).args
-    names_idx = {name: i for i, name in enumerate(names)}
-    args_order = []
-    kwargs_order = {}
-    for name, arg in kwargs.items():
-        if name in names:
-            kwargs_order[name] = names_idx[name]
-            names.remove(name)
-    for i, arg in enumerate(args):
-        if i < len(names):
-            args_order.append(names_idx[names[i]])
-    return args_order, kwargs_order
-
-
-def broadcast_args(args, kwargs, args_dim, kwargs_dim):
-    spatial = []
-    for arg, arg_dim in zip(
-        args + list(kwargs.values()), args_dim + list(kwargs_dim.values())
-    ):
-        if isinstance(arg, torch.Tensor) and arg_dim is not None:
-            arg_spatial = arg.shape[: arg.ndim - arg_dim]
-            if len(arg_spatial) > len(spatial):
-                spatial = [1] * (len(arg_spatial) - len(spatial)) + spatial
-            for j in range(len(arg_spatial)):
-                if spatial[-j] < arg_spatial[-j]:
-                    if spatial[-j] == 1:
-                        spatial[-j] = arg_spatial[-j]
-                    else:
-                        raise ValueError("Cannot broadcast arguments.")
-    for i, arg in enumerate(args):
-        if isinstance(arg, torch.Tensor) and args_dim[i] is not None:
-            args[i] = torch.broadcast_to(
-                arg, [*spatial, *arg.shape[arg.ndim - args_dim[i] :]]
-            )
-    for key, arg in kwargs.items():
-        if isinstance(arg, torch.Tensor) and kwargs_dim[key] is not None:
-            kwargs[key] = torch.broadcast_to(
-                arg, [*spatial, *arg.shape[arg.ndim - kwargs_dim[key] :]]
-            )
-    return args, kwargs, spatial
-
-
-@suppress_traceback
-def batched(*dims):
-    """
-    Decorator that allows a function to be called with batched arguments.
-    """
-
-    def decorator(func):
-        @wraps(func)
-        def wrapper(*args, device=torch.device("cpu"), **kwargs):
-            args = list(args)
-            # get arguments dimensions
-            args_order, kwargs_order = get_args_order(func, args, kwargs)
-            args_dim = [dims[i] for i in args_order]
-            kwargs_dim = {key: dims[i] for key, i in kwargs_order.items()}
-            # convert to torch tensor
-            device = get_device(args, kwargs) or device
-            for i, arg in enumerate(args):
-                if isinstance(arg, (Number, list, tuple)) and args_dim[i] is not None:
-                    args[i] = torch.tensor(arg, device=device)
-            for key, arg in kwargs.items():
-                if (
-                    isinstance(arg, (Number, list, tuple))
-                    and kwargs_dim[key] is not None
-                ):
-                    kwargs[key] = torch.tensor(arg, device=device)
-            # broadcast arguments
-            args, kwargs, spatial = broadcast_args(args, kwargs, args_dim, kwargs_dim)
-            for i, (arg, arg_dim) in enumerate(zip(args, args_dim)):
-                if isinstance(arg, torch.Tensor) and arg_dim is not None:
-                    args[i] = arg.reshape([-1, *arg.shape[arg.ndim - arg_dim :]])
-            for key, arg in kwargs.items():
-                if isinstance(arg, torch.Tensor) and kwargs_dim[key] is not None:
-                    kwargs[key] = arg.reshape(
-                        [-1, *arg.shape[arg.ndim - kwargs_dim[key] :]]
-                    )
-            # call function
-            results = func(*args, **kwargs)
-            type_results = type(results)
-            results = list(results) if isinstance(results, (tuple, list)) else [results]
-            # restore spatial dimensions
-            for i, result in enumerate(results):
-                results[i] = result.reshape([*spatial, *result.shape[1:]])
-            if type_results is tuple:
-                results = tuple(results)
-            elif type_results is list:
-                results = list(results)
-            else:
-                results = results[0]
-            return results
-
-        return wrapper
-
-    return decorator
-
-
-def image_uv(
-    height: int,
-    width: int,
-    left: int = None,
-    top: int = None,
-    right: int = None,
-    bottom: int = None,
-    device: torch.device = None,
-    dtype: torch.dtype = None,
-) -> torch.Tensor:
-    """
-    Get image space UV grid, ranging in [0, 1].
-
-    >>> image_uv(10, 10):
-    [[[0.05, 0.05], [0.15, 0.05], ..., [0.95, 0.05]],
-     [[0.05, 0.15], [0.15, 0.15], ..., [0.95, 0.15]],
-      ...             ...                  ...
-     [[0.05, 0.95], [0.15, 0.95], ..., [0.95, 0.95]]]
-
-    Args:
-        width (int): image width
-        height (int): image height
-
-    Returns:
-        torch.Tensor: shape (height, width, 2)
-    """
-    if left is None:
-        left = 0
-    if top is None:
-        top = 0
-    if right is None:
-        right = width
-    if bottom is None:
-        bottom = height
-    u = torch.linspace(
-        (left + 0.5) / width,
-        (right - 0.5) / width,
-        right - left,
-        device=device,
-        dtype=dtype,
-    )
-    v = torch.linspace(
-        (top + 0.5) / height,
-        (bottom - 0.5) / height,
-        bottom - top,
-        device=device,
-        dtype=dtype,
-    )
-    u, v = torch.meshgrid(u, v, indexing="xy")
-    uv = torch.stack([u, v], dim=-1)
-
-    return uv
-
-
-@batched(2, 1, 2, 2)
-def unproject_cv(
-    uv_coord: torch.Tensor,
-    depth: torch.Tensor = None,
-    extrinsics: torch.Tensor = None,
-    intrinsics: torch.Tensor = None,
-) -> torch.Tensor:
-    """
-    Unproject uv coordinates to 3D view space following the OpenCV convention
-
-    Args:
-        uv_coord (torch.Tensor): [..., N, 2] uv coordinates, value ranging in [0, 1].
-            The origin (0., 0.) is corresponding to the left & top
-        depth (torch.Tensor): [..., N] depth value
-        extrinsics (torch.Tensor): [..., 4, 4] extrinsics matrix
-        intrinsics (torch.Tensor): [..., 3, 3] intrinsics matrix
-
-    Returns:
-        points (torch.Tensor): [..., N, 3] 3d points
-    """
-    assert intrinsics is not None, "intrinsics matrix is required"
-    points = torch.cat([uv_coord, torch.ones_like(uv_coord[..., :1])], dim=-1)
-    points = points @ torch.inverse(intrinsics).transpose(-2, -1)
-    if depth is not None:
-        points = points * depth[..., None]
-    if extrinsics is not None:
-        points = torch.cat([points, torch.ones_like(points[..., :1])], dim=-1)
-        points = (points @ torch.inverse(extrinsics).transpose(-2, -1))[..., :3]
-    return points
-
-
-def depth_to_points(
-    depth: torch.Tensor, intrinsics: torch.Tensor, extrinsics: torch.Tensor = None
-):
-    height, width = depth.shape[-2:]
-    uv = image_uv(width=width, height=height, dtype=depth.dtype, device=depth.device)
-    pts = unproject_cv(
-        uv,
-        depth,
-        intrinsics=intrinsics[..., None, :, :],
-        extrinsics=extrinsics[..., None, :, :] if extrinsics is not None else None,
-    )
-
-    return pts
-
-
-@batched(0, 0, 0, 0, 0, 0)
-def intrinsics_from_focal_center(
-    fx: Union[float, torch.Tensor],
-    fy: Union[float, torch.Tensor],
-    cx: Union[float, torch.Tensor],
-    cy: Union[float, torch.Tensor],
-) -> torch.Tensor:
-    """
-    Get OpenCV intrinsics matrix
-
-    Args:
-        focal_x (float | torch.Tensor): focal length in x axis
-        focal_y (float | torch.Tensor): focal length in y axis
-        cx (float | torch.Tensor): principal point in x axis
-        cy (float | torch.Tensor): principal point in y axis
-
-    Returns:
-        (torch.Tensor): [..., 3, 3] OpenCV intrinsics matrix
-    """
-    N = fx.shape[0]
-    ret = torch.zeros((N, 3, 3), dtype=fx.dtype, device=fx.device)
-    zeros, ones = (
-        torch.zeros(N, dtype=fx.dtype, device=fx.device),
-        torch.ones(N, dtype=fx.dtype, device=fx.device),
-    )
-    ret = torch.stack(
-        [fx, zeros, cx, zeros, fy, cy, zeros, zeros, ones], dim=-1
-    ).unflatten(-1, (3, 3))
-    return ret

mapanything/models/external/moge/models/v1.py

@@ -1,595 +0,0 @@
-import importlib
-from numbers import Number
-from pathlib import Path
-from typing import Any, Dict, IO, List, Literal, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils
-import torch.utils.checkpoint
-import torch.version
-from huggingface_hub import hf_hub_download
-
-from mapanything.models.external.moge.models.utils import (
-    depth_to_points,
-    intrinsics_from_focal_center,
-    normalized_view_plane_uv,
-    recover_focal_shift,
-    wrap_module_with_gradient_checkpointing,
-)
-
-
-class ResidualConvBlock(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int = None,
-        hidden_channels: int = None,
-        padding_mode: str = "replicate",
-        activation: Literal["relu", "leaky_relu", "silu", "elu"] = "relu",
-        norm: Literal["group_norm", "layer_norm"] = "group_norm",
-    ):
-        super(ResidualConvBlock, self).__init__()
-        if out_channels is None:
-            out_channels = in_channels
-        if hidden_channels is None:
-            hidden_channels = in_channels
-
-        if activation == "relu":
-            activation_cls = lambda: nn.ReLU(inplace=True)  # noqa
-        elif activation == "leaky_relu":
-            activation_cls = lambda: nn.LeakyReLU(negative_slope=0.2, inplace=True)  # noqa
-        elif activation == "silu":
-            activation_cls = lambda: nn.SiLU(inplace=True)  # noqa
-        elif activation == "elu":
-            activation_cls = lambda: nn.ELU(inplace=True)  # noqa
-        else:
-            raise ValueError(f"Unsupported activation function: {activation}")
-
-        self.layers = nn.Sequential(
-            nn.GroupNorm(1, in_channels),
-            activation_cls(),
-            nn.Conv2d(
-                in_channels,
-                hidden_channels,
-                kernel_size=3,
-                padding=1,
-                padding_mode=padding_mode,
-            ),
-            nn.GroupNorm(
-                hidden_channels // 32 if norm == "group_norm" else 1, hidden_channels
-            ),
-            activation_cls(),
-            nn.Conv2d(
-                hidden_channels,
-                out_channels,
-                kernel_size=3,
-                padding=1,
-                padding_mode=padding_mode,
-            ),
-        )
-
-        self.skip_connection = (
-            nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)
-            if in_channels != out_channels
-            else nn.Identity()
-        )
-
-    def forward(self, x):
-        skip = self.skip_connection(x)
-        x = self.layers(x)
-        x = x + skip
-        return x
-
-
-class Head(nn.Module):
-    def __init__(
-        self,
-        num_features: int,
-        dim_in: int,
-        dim_out: List[int],
-        dim_proj: int = 512,
-        dim_upsample: List[int] = [256, 128, 128],
-        dim_times_res_block_hidden: int = 1,
-        num_res_blocks: int = 1,
-        res_block_norm: Literal["group_norm", "layer_norm"] = "group_norm",
-        last_res_blocks: int = 0,
-        last_conv_channels: int = 32,
-        last_conv_size: int = 1,
-    ):
-        super().__init__()
-
-        self.projects = nn.ModuleList(
-            [
-                nn.Conv2d(
-                    in_channels=dim_in,
-                    out_channels=dim_proj,
-                    kernel_size=1,
-                    stride=1,
-                    padding=0,
-                )
-                for _ in range(num_features)
-            ]
-        )
-
-        self.upsample_blocks = nn.ModuleList(
-            [
-                nn.Sequential(
-                    self._make_upsampler(in_ch + 2, out_ch),
-                    *(
-                        ResidualConvBlock(
-                            out_ch,
-                            out_ch,
-                            dim_times_res_block_hidden * out_ch,
-                            activation="relu",
-                            norm=res_block_norm,
-                        )
-                        for _ in range(num_res_blocks)
-                    ),
-                )
-                for in_ch, out_ch in zip([dim_proj] + dim_upsample[:-1], dim_upsample)
-            ]
-        )
-
-        self.output_block = nn.ModuleList(
-            [
-                self._make_output_block(
-                    dim_upsample[-1] + 2,
-                    dim_out_,
-                    dim_times_res_block_hidden,
-                    last_res_blocks,
-                    last_conv_channels,
-                    last_conv_size,
-                    res_block_norm,
-                )
-                for dim_out_ in dim_out
-            ]
-        )
-
-    def _make_upsampler(self, in_channels: int, out_channels: int):
-        upsampler = nn.Sequential(
-            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),
-            nn.Conv2d(
-                out_channels,
-                out_channels,
-                kernel_size=3,
-                stride=1,
-                padding=1,
-                padding_mode="replicate",
-            ),
-        )
-        upsampler[0].weight.data[:] = upsampler[0].weight.data[:, :, :1, :1]
-        return upsampler
-
-    def _make_output_block(
-        self,
-        dim_in: int,
-        dim_out: int,
-        dim_times_res_block_hidden: int,
-        last_res_blocks: int,
-        last_conv_channels: int,
-        last_conv_size: int,
-        res_block_norm: Literal["group_norm", "layer_norm"],
-    ):
-        return nn.Sequential(
-            nn.Conv2d(
-                dim_in,
-                last_conv_channels,
-                kernel_size=3,
-                stride=1,
-                padding=1,
-                padding_mode="replicate",
-            ),
-            *(
-                ResidualConvBlock(
-                    last_conv_channels,
-                    last_conv_channels,
-                    dim_times_res_block_hidden * last_conv_channels,
-                    activation="relu",
-                    norm=res_block_norm,
-                )
-                for _ in range(last_res_blocks)
-            ),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(
-                last_conv_channels,
-                dim_out,
-                kernel_size=last_conv_size,
-                stride=1,
-                padding=last_conv_size // 2,
-                padding_mode="replicate",
-            ),
-        )
-
-    def forward(self, hidden_states: torch.Tensor, image: torch.Tensor):
-        img_h, img_w = image.shape[-2:]
-        patch_h, patch_w = img_h // 14, img_w // 14
-
-        # Process the hidden states
-        x = torch.stack(
-            [
-                proj(
-                    feat.permute(0, 2, 1).unflatten(2, (patch_h, patch_w)).contiguous()
-                )
-                for proj, (feat, clstoken) in zip(self.projects, hidden_states)
-            ],
-            dim=1,
-        ).sum(dim=1)
-
-        # Upsample stage
-        # (patch_h, patch_w) -> (patch_h * 2, patch_w * 2) -> (patch_h * 4, patch_w * 4) -> (patch_h * 8, patch_w * 8)
-        for i, block in enumerate(self.upsample_blocks):
-            # UV coordinates is for awareness of image aspect ratio
-            uv = normalized_view_plane_uv(
-                width=x.shape[-1],
-                height=x.shape[-2],
-                aspect_ratio=img_w / img_h,
-                dtype=x.dtype,
-                device=x.device,
-            )
-            uv = uv.permute(2, 0, 1).unsqueeze(0).expand(x.shape[0], -1, -1, -1)
-            x = torch.cat([x, uv], dim=1)
-            for layer in block:
-                x = torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False)
-
-        # (patch_h * 8, patch_w * 8) -> (img_h, img_w)
-        x = F.interpolate(x, (img_h, img_w), mode="bilinear", align_corners=False)
-        uv = normalized_view_plane_uv(
-            width=x.shape[-1],
-            height=x.shape[-2],
-            aspect_ratio=img_w / img_h,
-            dtype=x.dtype,
-            device=x.device,
-        )
-        uv = uv.permute(2, 0, 1).unsqueeze(0).expand(x.shape[0], -1, -1, -1)
-        x = torch.cat([x, uv], dim=1)
-
-        if isinstance(self.output_block, nn.ModuleList):
-            output = [
-                torch.utils.checkpoint.checkpoint(block, x, use_reentrant=False)
-                for block in self.output_block
-            ]
-        else:
-            output = torch.utils.checkpoint.checkpoint(
-                self.output_block, x, use_reentrant=False
-            )
-
-        return output
-
-
-class MoGeModel(nn.Module):
-    image_mean: torch.Tensor
-    image_std: torch.Tensor
-
-    def __init__(
-        self,
-        encoder: str = "dinov2_vitb14",
-        intermediate_layers: Union[int, List[int]] = 4,
-        dim_proj: int = 512,
-        dim_upsample: List[int] = [256, 128, 128],
-        dim_times_res_block_hidden: int = 1,
-        num_res_blocks: int = 1,
-        remap_output: Literal[
-            False, True, "linear", "sinh", "exp", "sinh_exp"
-        ] = "linear",
-        res_block_norm: Literal["group_norm", "layer_norm"] = "group_norm",
-        num_tokens_range: Tuple[Number, Number] = [1200, 2500],
-        last_res_blocks: int = 0,
-        last_conv_channels: int = 32,
-        last_conv_size: int = 1,
-        mask_threshold: float = 0.5,
-        **deprecated_kwargs,
-    ):
-        super(MoGeModel, self).__init__()
-
-        if deprecated_kwargs:
-            # Process legacy arguments
-            if "trained_area_range" in deprecated_kwargs:
-                num_tokens_range = [
-                    deprecated_kwargs["trained_area_range"][0] // 14**2,
-                    deprecated_kwargs["trained_area_range"][1] // 14**2,
-                ]
-                del deprecated_kwargs["trained_area_range"]
-            # warnings.warn(
-            #     f"The following deprecated/invalid arguments are ignored: {deprecated_kwargs}"
-            # )
-
-        self.encoder = encoder
-        self.remap_output = remap_output
-        self.intermediate_layers = intermediate_layers
-        self.num_tokens_range = num_tokens_range
-        self.mask_threshold = mask_threshold
-
-        # NOTE: We have copied the DINOv2 code in torchhub to this repository.
-        # Minimal modifications have been made: removing irrelevant code, unnecessary warnings and fixing importing issues.
-        hub_loader = getattr(
-            importlib.import_module(
-                "mapanything.models.external.dinov2.hub.backbones", __package__
-            ),
-            encoder,
-        )
-        self.backbone = hub_loader(pretrained=False)
-        dim_feature = self.backbone.blocks[0].attn.qkv.in_features
-
-        self.head = Head(
-            num_features=intermediate_layers
-            if isinstance(intermediate_layers, int)
-            else len(intermediate_layers),
-            dim_in=dim_feature,
-            dim_out=[3, 1],
-            dim_proj=dim_proj,
-            dim_upsample=dim_upsample,
-            dim_times_res_block_hidden=dim_times_res_block_hidden,
-            num_res_blocks=num_res_blocks,
-            res_block_norm=res_block_norm,
-            last_res_blocks=last_res_blocks,
-            last_conv_channels=last_conv_channels,
-            last_conv_size=last_conv_size,
-        )
-
-        image_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
-        image_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
-
-        self.register_buffer("image_mean", image_mean)
-        self.register_buffer("image_std", image_std)
-
-    @property
-    def device(self) -> torch.device:
-        return next(self.parameters()).device
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return next(self.parameters()).dtype
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: Union[str, Path, IO[bytes]],
-        model_kwargs: Optional[Dict[str, Any]] = None,
-        **hf_kwargs,
-    ) -> "MoGeModel":
-        """
-        Load a model from a checkpoint file.
-
-        ### Parameters:
-        - `pretrained_model_name_or_path`: path to the checkpoint file or repo id.
-        - `model_kwargs`: additional keyword arguments to override the parameters in the checkpoint.
-        - `hf_kwargs`: additional keyword arguments to pass to the `hf_hub_download` function. Ignored if `pretrained_model_name_or_path` is a local path.
-
-        ### Returns:
-        - A new instance of `MoGe` with the parameters loaded from the checkpoint.
-        """
-        if Path(pretrained_model_name_or_path).exists():
-            checkpoint = torch.load(
-                pretrained_model_name_or_path, map_location="cpu", weights_only=True
-            )
-        else:
-            cached_checkpoint_path = hf_hub_download(
-                repo_id=pretrained_model_name_or_path,
-                repo_type="model",
-                filename="model.pt",
-                **hf_kwargs,
-            )
-            checkpoint = torch.load(
-                cached_checkpoint_path, map_location="cpu", weights_only=True
-            )
-        model_config = checkpoint["model_config"]
-        if model_kwargs is not None:
-            model_config.update(model_kwargs)
-        model = cls(**model_config)
-        model.load_state_dict(checkpoint["model"])
-        return model
-
-    def init_weights(self):
-        "Load the backbone with pretrained dinov2 weights from torch hub"
-        state_dict = torch.hub.load(
-            "facebookresearch/dinov2", self.encoder, pretrained=True
-        ).state_dict()
-        self.backbone.load_state_dict(state_dict)
-
-    def enable_gradient_checkpointing(self):
-        for i in range(len(self.backbone.blocks)):
-            self.backbone.blocks[i] = wrap_module_with_gradient_checkpointing(
-                self.backbone.blocks[i]
-            )
-
-    def _remap_points(self, points: torch.Tensor) -> torch.Tensor:
-        if self.remap_output == "linear":
-            pass
-        elif self.remap_output == "sinh":
-            points = torch.sinh(points)
-        elif self.remap_output == "exp":
-            xy, z = points.split([2, 1], dim=-1)
-            z = torch.exp(z)
-            points = torch.cat([xy * z, z], dim=-1)
-        elif self.remap_output == "sinh_exp":
-            xy, z = points.split([2, 1], dim=-1)
-            points = torch.cat([torch.sinh(xy), torch.exp(z)], dim=-1)
-        else:
-            raise ValueError(f"Invalid remap output type: {self.remap_output}")
-        return points
-
-    def forward(self, image: torch.Tensor, num_tokens: int) -> Dict[str, torch.Tensor]:
-        original_height, original_width = image.shape[-2:]
-
-        # Resize to expected resolution defined by num_tokens
-        resize_factor = (
-            (num_tokens * 14**2) / (original_height * original_width)
-        ) ** 0.5
-        resized_width, resized_height = (
-            int(original_width * resize_factor),
-            int(original_height * resize_factor),
-        )
-        image = F.interpolate(
-            image,
-            (resized_height, resized_width),
-            mode="bicubic",
-            align_corners=False,
-            antialias=True,
-        )
-
-        # Apply image transformation for DINOv2
-        image = (image - self.image_mean) / self.image_std
-        image_14 = F.interpolate(
-            image,
-            (resized_height // 14 * 14, resized_width // 14 * 14),
-            mode="bilinear",
-            align_corners=False,
-            antialias=True,
-        )
-
-        # Get intermediate layers from the backbone
-        features = self.backbone.get_intermediate_layers(
-            image_14, self.intermediate_layers, return_class_token=True
-        )
-
-        # Predict points (and mask)
-        output = self.head(features, image)
-        points, mask = output
-
-        # Make sure fp32 precision for output
-        with torch.autocast(device_type=image.device.type, dtype=torch.float32):
-            # Resize to original resolution
-            points = F.interpolate(
-                points,
-                (original_height, original_width),
-                mode="bilinear",
-                align_corners=False,
-                antialias=False,
-            )
-            mask = F.interpolate(
-                mask,
-                (original_height, original_width),
-                mode="bilinear",
-                align_corners=False,
-                antialias=False,
-            )
-
-            # Post-process points and mask
-            points, mask = points.permute(0, 2, 3, 1), mask.squeeze(1)
-            points = self._remap_points(
-                points
-            )  # slightly improves the performance in case of very large output values
-
-        return_dict = {"points": points, "mask": mask}
-        return return_dict
-
-    # @torch.inference_mode()
-    def infer(
-        self,
-        image: torch.Tensor,
-        fov_x: Union[Number, torch.Tensor] = None,
-        resolution_level: int = 9,
-        num_tokens: int = None,
-        apply_mask: bool = True,
-        force_projection: bool = True,
-        use_fp16: bool = True,
-    ) -> Dict[str, torch.Tensor]:
-        """
-        User-friendly inference function
-
-        ### Parameters
-        - `image`: input image tensor of shape (B, 3, H, W) or (3, H, W)\
-        - `fov_x`: the horizontal camera FoV in degrees. If None, it will be inferred from the predicted point map. Default: None
-        - `resolution_level`: An integer [0-9] for the resolution level for inference.
-            The higher, the finer details will be captured, but slower. Defaults to 9. Note that it is irrelevant to the output size, which is always the same as the input size.
-            `resolution_level` actually controls `num_tokens`. See `num_tokens` for more details.
-        - `num_tokens`: number of tokens used for inference. A integer in the (suggested) range of `[1200, 2500]`.
-            `resolution_level` will be ignored if `num_tokens` is provided. Default: None
-        - `apply_mask`: if True, the output point map will be masked using the predicted mask. Default: True
-        - `force_projection`: if True, the output point map will be recomputed to match the projection constraint. Default: True
-        - `use_fp16`: if True, use mixed precision to speed up inference. Default: True
-
-        ### Returns
-
-        A dictionary containing the following keys:
-        - `points`: output tensor of shape (B, H, W, 3) or (H, W, 3).
-        - `depth`: tensor of shape (B, H, W) or (H, W) containing the depth map.
-        - `intrinsics`: tensor of shape (B, 3, 3) or (3, 3) containing the camera intrinsics.
-        """
-        if image.dim() == 3:
-            omit_batch_dim = True
-            image = image.unsqueeze(0)
-        else:
-            omit_batch_dim = False
-        image = image.to(dtype=self.dtype, device=self.device)
-
-        original_height, original_width = image.shape[-2:]
-        aspect_ratio = original_width / original_height
-
-        if num_tokens is None:
-            min_tokens, max_tokens = self.num_tokens_range
-            num_tokens = int(
-                min_tokens + (resolution_level / 9) * (max_tokens - min_tokens)
-            )
-
-        with torch.autocast(
-            device_type=self.device.type,
-            dtype=torch.float16,
-            enabled=use_fp16 and self.dtype != torch.float16,
-        ):
-            output = self.forward(image, num_tokens)
-        points, mask = output["points"], output["mask"]
-
-        # Always process the output in fp32 precision
-        with torch.autocast(device_type=self.device.type, dtype=torch.float32):
-            points, mask, fov_x = map(
-                lambda x: x.float() if isinstance(x, torch.Tensor) else x,
-                [points, mask, fov_x],
-            )
-
-            mask_binary = mask > self.mask_threshold
-
-            # Get camera-space point map. (Focal here is the focal length relative to half the image diagonal)
-            if fov_x is None:
-                focal, shift = recover_focal_shift(points, mask_binary)
-            else:
-                focal = (
-                    aspect_ratio
-                    / (1 + aspect_ratio**2) ** 0.5
-                    / torch.tan(
-                        torch.deg2rad(
-                            torch.as_tensor(
-                                fov_x, device=points.device, dtype=points.dtype
-                            )
-                            / 2
-                        )
-                    )
-                )
-                if focal.ndim == 0:
-                    focal = focal[None].expand(points.shape[0])
-                _, shift = recover_focal_shift(points, mask_binary, focal=focal)
-            fx = focal / 2 * (1 + aspect_ratio**2) ** 0.5 / aspect_ratio
-            fy = focal / 2 * (1 + aspect_ratio**2) ** 0.5
-            intrinsics = intrinsics_from_focal_center(fx, fy, 0.5, 0.5)
-            depth = points[..., 2] + shift[..., None, None]
-
-            # If projection constraint is forced, recompute the point map using the actual depth map
-            if force_projection:
-                points = depth_to_points(depth, intrinsics=intrinsics)
-            else:
-                points = (
-                    points
-                    + torch.stack(
-                        [torch.zeros_like(shift), torch.zeros_like(shift), shift],
-                        dim=-1,
-                    )[..., None, None, :]
-                )
-
-            # Apply mask if needed
-            if apply_mask:
-                points = torch.where(mask_binary[..., None], points, torch.inf)
-                depth = torch.where(mask_binary, depth, torch.inf)
-
-        return_dict = {
-            "points": points,
-            "intrinsics": intrinsics,
-            "depth": depth,
-            "mask": mask_binary,
-        }
-
-        if omit_batch_dim:
-            return_dict = {k: v.squeeze(0) for k, v in return_dict.items()}
-
-        return return_dict

mapanything/models/external/moge/models/v2.py

@@ -1,379 +0,0 @@
-import warnings
-from numbers import Number
-from pathlib import Path
-from typing import Any, Dict, IO, List, Literal, Optional, Union
-
-import torch
-import torch.amp
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils
-import torch.utils.checkpoint
-import torch.version
-from huggingface_hub import hf_hub_download
-
-from mapanything.models.external.moge.models.modules import (
-    ConvStack,
-    DINOv2Encoder,
-    MLP,
-)
-from mapanything.models.external.moge.models.utils import (
-    depth_to_points,
-    intrinsics_from_focal_center,
-    normalized_view_plane_uv,
-    recover_focal_shift,
-)
-
-
-class MoGeModel(nn.Module):
-    encoder: DINOv2Encoder
-    neck: ConvStack
-    points_head: ConvStack
-    mask_head: ConvStack
-    scale_head: MLP
-    onnx_compatible_mode: bool
-
-    def __init__(
-        self,
-        encoder: Dict[str, Any],
-        neck: Dict[str, Any],
-        points_head: Dict[str, Any] = None,
-        mask_head: Dict[str, Any] = None,
-        normal_head: Dict[str, Any] = None,
-        scale_head: Dict[str, Any] = None,
-        remap_output: Literal["linear", "sinh", "exp", "sinh_exp"] = "linear",
-        num_tokens_range: List[int] = [1200, 3600],
-        **deprecated_kwargs,
-    ):
-        super(MoGeModel, self).__init__()
-        if deprecated_kwargs:
-            warnings.warn(
-                f"The following deprecated/invalid arguments are ignored: {deprecated_kwargs}"
-            )
-
-        self.remap_output = remap_output
-        self.num_tokens_range = num_tokens_range
-
-        self.encoder = DINOv2Encoder(**encoder)
-        self.neck = ConvStack(**neck)
-        if points_head is not None:
-            self.points_head = ConvStack(**points_head)
-        if mask_head is not None:
-            self.mask_head = ConvStack(**mask_head)
-        if normal_head is not None:
-            self.normal_head = ConvStack(**normal_head)
-        if scale_head is not None:
-            self.scale_head = MLP(**scale_head)
-
-    @property
-    def device(self) -> torch.device:
-        return next(self.parameters()).device
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return next(self.parameters()).dtype
-
-    @property
-    def onnx_compatible_mode(self) -> bool:
-        return getattr(self, "_onnx_compatible_mode", False)
-
-    @onnx_compatible_mode.setter
-    def onnx_compatible_mode(self, value: bool):
-        self._onnx_compatible_mode = value
-        self.encoder.onnx_compatible_mode = value
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: Union[str, Path, IO[bytes]],
-        model_kwargs: Optional[Dict[str, Any]] = None,
-        **hf_kwargs,
-    ) -> "MoGeModel":
-        """
-        Load a model from a checkpoint file.
-
-        ### Parameters:
-        - `pretrained_model_name_or_path`: path to the checkpoint file or repo id.
-        - `compiled`
-        - `model_kwargs`: additional keyword arguments to override the parameters in the checkpoint.
-        - `hf_kwargs`: additional keyword arguments to pass to the `hf_hub_download` function. Ignored if `pretrained_model_name_or_path` is a local path.
-
-        ### Returns:
-        - A new instance of `MoGe` with the parameters loaded from the checkpoint.
-        """
-        if Path(pretrained_model_name_or_path).exists():
-            checkpoint_path = pretrained_model_name_or_path
-        else:
-            checkpoint_path = hf_hub_download(
-                repo_id=pretrained_model_name_or_path,
-                repo_type="model",
-                filename="model.pt",
-                **hf_kwargs,
-            )
-        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
-
-        model_config = checkpoint["model_config"]
-        if model_kwargs is not None:
-            model_config.update(model_kwargs)
-        model = cls(**model_config)
-        model.load_state_dict(checkpoint["model"], strict=False)
-
-        return model
-
-    def init_weights(self):
-        self.encoder.init_weights()
-
-    def enable_gradient_checkpointing(self):
-        self.encoder.enable_gradient_checkpointing()
-        self.neck.enable_gradient_checkpointing()
-        for head in ["points_head", "normal_head", "mask_head"]:
-            if hasattr(self, head):
-                getattr(self, head).enable_gradient_checkpointing()
-
-    def enable_pytorch_native_sdpa(self):
-        self.encoder.enable_pytorch_native_sdpa()
-
-    def _remap_points(self, points: torch.Tensor) -> torch.Tensor:
-        if self.remap_output == "linear":
-            pass
-        elif self.remap_output == "sinh":
-            points = torch.sinh(points)
-        elif self.remap_output == "exp":
-            xy, z = points.split([2, 1], dim=-1)
-            z = torch.exp(z)
-            points = torch.cat([xy * z, z], dim=-1)
-        elif self.remap_output == "sinh_exp":
-            xy, z = points.split([2, 1], dim=-1)
-            points = torch.cat([torch.sinh(xy), torch.exp(z)], dim=-1)
-        else:
-            raise ValueError(f"Invalid remap output type: {self.remap_output}")
-        return points
-
-    def forward(self, image: torch.Tensor, num_tokens: int) -> Dict[str, torch.Tensor]:
-        batch_size, _, img_h, img_w = image.shape
-        device, dtype = image.device, image.dtype
-
-        aspect_ratio = img_w / img_h
-        base_h, base_w = (
-            int((num_tokens / aspect_ratio) ** 0.5),
-            int((num_tokens * aspect_ratio) ** 0.5),
-        )
-        num_tokens = base_h * base_w
-
-        # Backbones encoding
-        features, cls_token = self.encoder(
-            image, base_h, base_w, return_class_token=True
-        )
-        features = [features, None, None, None, None]
-
-        # Concat UVs for aspect ratio input
-        for level in range(5):
-            uv = normalized_view_plane_uv(
-                width=base_w * 2**level,
-                height=base_h * 2**level,
-                aspect_ratio=aspect_ratio,
-                dtype=dtype,
-                device=device,
-            )
-            uv = uv.permute(2, 0, 1).unsqueeze(0).expand(batch_size, -1, -1, -1)
-            if features[level] is None:
-                features[level] = uv
-            else:
-                features[level] = torch.concat([features[level], uv], dim=1)
-
-        # Shared neck
-        features = self.neck(features)
-
-        # Heads decoding
-        points, normal, mask = (
-            getattr(self, head)(features)[-1] if hasattr(self, head) else None
-            for head in ["points_head", "normal_head", "mask_head"]
-        )
-        metric_scale = (
-            self.scale_head(cls_token) if hasattr(self, "scale_head") else None
-        )
-
-        # Resize
-        points, normal, mask = (
-            F.interpolate(
-                v, (img_h, img_w), mode="bilinear", align_corners=False, antialias=False
-            )
-            if v is not None
-            else None
-            for v in [points, normal, mask]
-        )
-
-        # Remap output
-        if points is not None:
-            points = points.permute(0, 2, 3, 1)
-            points = self._remap_points(
-                points
-            )  # slightly improves the performance in case of very large output values
-        if normal is not None:
-            normal = normal.permute(0, 2, 3, 1)
-            normal = F.normalize(normal, dim=-1)
-        if mask is not None:
-            mask = mask.squeeze(1).sigmoid()
-        if metric_scale is not None:
-            metric_scale = metric_scale.squeeze(1).exp()
-
-        return_dict = {
-            "points": points,
-            "normal": normal,
-            "mask": mask,
-            "metric_scale": metric_scale,
-        }
-        return_dict = {k: v for k, v in return_dict.items() if v is not None}
-
-        return return_dict
-
-    # @torch.inference_mode()
-    def infer(
-        self,
-        image: torch.Tensor,
-        num_tokens: int = None,
-        resolution_level: int = 9,
-        force_projection: bool = True,
-        apply_mask: Literal[False, True, "blend"] = True,
-        fov_x: Optional[Union[Number, torch.Tensor]] = None,
-        use_fp16: bool = True,
-    ) -> Dict[str, torch.Tensor]:
-        """
-        User-friendly inference function
-
-        ### Parameters
-        - `image`: input image tensor of shape (B, 3, H, W) or (3, H, W)
-        - `num_tokens`: the number of base ViT tokens to use for inference, `'least'` or `'most'` or an integer. Suggested range: 1200 ~ 2500.
-            More tokens will result in significantly higher accuracy and finer details, but slower inference time. Default: `'most'`.
-        - `force_projection`: if True, the output point map will be computed using the actual depth map. Default: True
-        - `apply_mask`: if True, the output point map will be masked using the predicted mask. Default: True
-        - `fov_x`: the horizontal camera FoV in degrees. If None, it will be inferred from the predicted point map. Default: None
-        - `use_fp16`: if True, use mixed precision to speed up inference. Default: True
-
-        ### Returns
-
-        A dictionary containing the following keys:
-        - `points`: output tensor of shape (B, H, W, 3) or (H, W, 3).
-        - `depth`: tensor of shape (B, H, W) or (H, W) containing the depth map.
-        - `intrinsics`: tensor of shape (B, 3, 3) or (3, 3) containing the camera intrinsics.
-        """
-        if image.dim() == 3:
-            omit_batch_dim = True
-            image = image.unsqueeze(0)
-        else:
-            omit_batch_dim = False
-        image = image.to(dtype=self.dtype, device=self.device)
-
-        original_height, original_width = image.shape[-2:]
-        aspect_ratio = original_width / original_height
-
-        # Determine the number of base tokens to use
-        if num_tokens is None:
-            min_tokens, max_tokens = self.num_tokens_range
-            num_tokens = int(
-                min_tokens + (resolution_level / 9) * (max_tokens - min_tokens)
-            )
-
-        # Forward pass
-        with torch.autocast(
-            device_type=self.device.type,
-            dtype=torch.float16,
-            enabled=use_fp16 and self.dtype != torch.float16,
-        ):
-            output = self.forward(image, num_tokens=num_tokens)
-        points, normal, mask, metric_scale = (
-            output.get(k, None) for k in ["points", "normal", "mask", "metric_scale"]
-        )
-
-        # Always process the output in fp32 precision
-        points, normal, mask, metric_scale, fov_x = map(
-            lambda x: x.float() if isinstance(x, torch.Tensor) else x,
-            [points, normal, mask, metric_scale, fov_x],
-        )
-        with torch.autocast(device_type=self.device.type, dtype=torch.float32):
-            if mask is not None:
-                mask_binary = mask > 0.5
-            else:
-                mask_binary = None
-
-            if points is not None:
-                # Convert affine point map to camera-space. Recover depth and intrinsics from point map.
-                # NOTE: Focal here is the focal length relative to half the image diagonal
-                if fov_x is None:
-                    # Recover focal and shift from predicted point map
-                    focal, shift = recover_focal_shift(points, mask_binary)
-                else:
-                    # Focal is known, recover shift only
-                    focal = (
-                        aspect_ratio
-                        / (1 + aspect_ratio**2) ** 0.5
-                        / torch.tan(
-                            torch.deg2rad(
-                                torch.as_tensor(
-                                    fov_x, device=points.device, dtype=points.dtype
-                                )
-                                / 2
-                            )
-                        )
-                    )
-                    if focal.ndim == 0:
-                        focal = focal[None].expand(points.shape[0])
-                    _, shift = recover_focal_shift(points, mask_binary, focal=focal)
-                fx, fy = (
-                    focal / 2 * (1 + aspect_ratio**2) ** 0.5 / aspect_ratio,
-                    focal / 2 * (1 + aspect_ratio**2) ** 0.5,
-                )
-                intrinsics = intrinsics_from_focal_center(fx, fy, 0.5, 0.5)
-                points[..., 2] += shift[..., None, None]
-                if mask_binary is not None:
-                    mask_binary &= (
-                        points[..., 2] > 0
-                    )  # in case depth is contains negative values (which should never happen in practice)
-                depth = points[..., 2].clone()
-            else:
-                depth, intrinsics = None, None
-
-            # If projection constraint is forced, recompute the point map using the actual depth map & intrinsics
-            if force_projection and depth is not None:
-                points = depth_to_points(depth, intrinsics=intrinsics)
-
-            # Apply metric scale
-            if metric_scale is not None:
-                if points is not None:
-                    points *= metric_scale[:, None, None, None]
-                if depth is not None:
-                    depth *= metric_scale[:, None, None]
-
-            # Apply mask
-            if apply_mask and mask_binary is not None:
-                points = (
-                    torch.where(mask_binary[..., None], points, torch.inf)
-                    if points is not None
-                    else None
-                )
-                depth = (
-                    torch.where(mask_binary, depth, torch.inf)
-                    if depth is not None
-                    else None
-                )
-                normal = (
-                    torch.where(
-                        mask_binary[..., None], normal, torch.zeros_like(normal)
-                    )
-                    if normal is not None
-                    else None
-                )
-
-        return_dict = {
-            "points": points,
-            "intrinsics": intrinsics,
-            "depth": depth,
-            "mask": mask_binary,
-            "normal": normal,
-        }
-        return_dict = {k: v for k, v in return_dict.items() if v is not None}
-
-        if omit_batch_dim:
-            return_dict = {k: v.squeeze(0) for k, v in return_dict.items()}
-
-        return return_dict

mapanything/models/external/must3r/__init__.py

@@ -1,283 +0,0 @@
-"""
-Inference wrapper for MUSt3R
-"""
-
-import datetime
-import os
-
-import numpy as np
-import torch
-from dust3r.viz import rgb
-from must3r.demo.inference import SceneState
-from must3r.engine.inference import inference_multi_ar, postprocess
-from must3r.model import get_pointmaps_activation, load_model
-
-from mapanything.models.external.vggt.utils.rotation import mat_to_quat
-
-
-def must3r_inference(
-    views,
-    filelist,
-    model,
-    retrieval,
-    device,
-    amp,
-    num_mem_images,
-    max_bs,
-    init_num_images=2,
-    batch_num_views=1,
-    render_once=False,
-    is_sequence=False,
-    viser_server=None,
-    num_refinements_iterations=2,
-    verbose=True,
-):
-    if amp == "fp16":
-        dtype = torch.float16
-    elif amp == "bf16":
-        assert torch.cuda.is_bf16_supported()
-        dtype = torch.bfloat16
-    else:
-        assert not amp
-        dtype = torch.float32
-
-    max_bs = None if max_bs == 0 else max_bs
-    encoder, decoder = model
-    pointmaps_activation = get_pointmaps_activation(decoder, verbose=verbose)
-
-    def post_process_function(x):
-        return postprocess(
-            x, pointmaps_activation=pointmaps_activation, compute_cam=True
-        )
-
-    if verbose:
-        print("loading images")
-    time_start = datetime.datetime.now()
-    nimgs = len(views)
-
-    ellapsed = datetime.datetime.now() - time_start
-    if verbose:
-        print(f"loaded in {ellapsed}")
-        print("running inference")
-    time_start = datetime.datetime.now()
-    if viser_server is not None:
-        viser_server.reset(nimgs)
-
-    imgs = [b["img"].to("cpu") for b in views]
-    true_shape = [torch.from_numpy(b["true_shape"]).to("cpu") for b in views]
-    true_shape = torch.stack(true_shape, dim=0)
-    nimgs = true_shape.shape[0]
-
-    # Use all images as keyframes
-    keyframes = np.linspace(0, len(imgs) - 1, num_mem_images, dtype=int).tolist()
-    encoder_precomputed_features = None
-
-    not_keyframes = sorted(set(range(nimgs)).difference(set(keyframes)))
-    assert (len(keyframes) + len(not_keyframes)) == nimgs
-    # reorder images
-    views = [views[i] for i in keyframes] + [views[i] for i in not_keyframes]
-    imgs = [b["img"].to(device) for b in views]
-    true_shape = [torch.from_numpy(b["true_shape"]).to(device) for b in views]
-    filenames = [filelist[i] for i in keyframes + not_keyframes]
-    img_ids = [torch.tensor(v) for v in keyframes + not_keyframes]
-
-    if encoder_precomputed_features is not None:
-        x_start, pos_start = encoder_precomputed_features
-        x = [x_start[i] for i in keyframes] + [x_start[i] for i in not_keyframes]
-        pos = [pos_start[i] for i in keyframes] + [pos_start[i] for i in not_keyframes]
-        encoder_precomputed_features = (x, pos)
-
-    mem_batches = [init_num_images]
-    while (sum_b := sum(mem_batches)) != max(num_mem_images, init_num_images):
-        size_b = min(batch_num_views, num_mem_images - sum_b)
-        mem_batches.append(size_b)
-
-    if render_once:
-        to_render = list(range(num_mem_images, nimgs))
-    else:
-        to_render = None
-
-    with torch.autocast("cuda", dtype=dtype):
-        x_out_0, x_out = inference_multi_ar(
-            encoder,
-            decoder,
-            imgs,
-            img_ids,
-            true_shape,
-            mem_batches,
-            max_bs=max_bs,
-            verbose=verbose,
-            to_render=to_render,
-            encoder_precomputed_features=encoder_precomputed_features,
-            device=device,
-            preserve_gpu_mem=True,
-            post_process_function=post_process_function,
-            viser_server=viser_server,
-            num_refinements_iterations=num_refinements_iterations,
-        )
-    if to_render is not None:
-        x_out = x_out_0 + x_out
-
-    ellapsed = datetime.datetime.now() - time_start
-    if verbose:
-        print(f"inference in {ellapsed}")
-        try:
-            print(str(int(torch.cuda.max_memory_reserved(device) / (1024**2))) + " MB")
-        except Exception:
-            pass
-
-    if viser_server is not None:
-        viser_server.reset_cam_visility()
-        viser_server.send_message("Finished")
-
-    if verbose:
-        print("preparing pointcloud")
-    time_start = datetime.datetime.now()
-    focals = []
-    cams2world = []
-    for i in range(nimgs):
-        focals.append(float(x_out[i]["focal"].cpu()))
-        cams2world.append(x_out[i]["c2w"].cpu())
-
-    # x_out to cpu
-    for i in range(len(x_out)):
-        for k in x_out[i].keys():
-            x_out[i][k] = x_out[i][k].cpu()
-
-    rgbimg = [rgb(imgs[i], true_shape[i]) for i in range(nimgs)]
-    scene = SceneState(x_out, rgbimg, true_shape, focals, cams2world, filenames)
-
-    ellapsed = datetime.datetime.now() - time_start
-    if verbose:
-        print(f"pointcloud prepared in {ellapsed}")
-
-    return scene
-
-
-class MUSt3RWrapper(torch.nn.Module):
-    def __init__(
-        self,
-        name,
-        ckpt_path,
-        retrieval_ckpt_path,
-        img_size=512,
-        amp="bf16",
-        max_bs=1,
-        **kwargs,
-    ):
-        super().__init__()
-        self.name = name
-        self.ckpt_path = ckpt_path
-        self.retrieval_ckpt_path = retrieval_ckpt_path
-        self.amp = amp
-        self.max_bs = max_bs
-
-        # Init the model and load the checkpoint
-        self.model = load_model(self.ckpt_path, img_size=512)
-
-    def forward(self, views):
-        """
-        Forward pass wrapper for MUSt3R.
-
-        Assumption:
-        - The batch size of input views is 1.
-
-        Args:
-            views (List[dict]): List of dictionaries containing the input views' images and instance information.
-                                Each dictionary should contain the following keys, where B is the batch size and is 1:
-                                    "img" (tensor): Image tensor of shape (B, C, H, W).
-                                    "data_norm_type" (list): ["dust3r"]
-                                    "label" (list): ["scene_name"]
-                                    "instance" (list): ["image_name"]
-
-        Returns:
-            List[dict]: A list containing the final outputs for the input views.
-        """
-        # Check the batch size of input views
-        batch_size_per_view, _, height, width = views[0]["img"].shape
-        device = views[0]["img"].device
-        num_views = len(views)
-        assert batch_size_per_view == 1, (
-            f"Batch size of input views should be 1, but got {batch_size_per_view}."
-        )
-
-        # Check the data norm type
-        data_norm_type = views[0]["data_norm_type"][0]
-        assert data_norm_type == "dust3r", (
-            "MUSt3R expects a normalized image with the DUSt3R normalization scheme applied"
-        )
-
-        # Convert the input views to the expected input format
-        images = []
-        image_paths = []
-        for view in views:
-            images.append(
-                dict(
-                    img=view["img"][0].cpu(),
-                    idx=len(images),
-                    instance=str(len(images)),
-                    true_shape=np.int32([view["img"].shape[-2], view["img"].shape[-1]]),
-                )
-            )
-            view_name = os.path.join(view["label"][0], view["instance"][0])
-            image_paths.append(view_name)
-
-        # Run MUSt3R inference
-        scene = must3r_inference(
-            images,
-            image_paths,
-            self.model,
-            self.retrieval_ckpt_path,
-            device,
-            self.amp,
-            num_views,
-            self.max_bs,
-            verbose=False,
-        )
-
-        # Make sure scene is not None
-        if scene is None:
-            raise RuntimeError("MUSt3R failed.")
-
-        # Get the predictions
-        predictions = scene.x_out
-
-        # Convert the output to the MapAnything format
-        with torch.autocast("cuda", enabled=False):
-            res = []
-            for view_idx in range(num_views):
-                # Get the current view predictions
-                curr_view_prediction = predictions[view_idx]
-                curr_view_conf = curr_view_prediction["conf"]
-                curr_view_pose = curr_view_prediction["c2w"].unsqueeze(0)
-
-                # Convert the pose to quaternions and translation
-                curr_view_cam_translations = curr_view_pose[..., :3, 3]
-                curr_view_cam_quats = mat_to_quat(curr_view_pose[..., :3, :3])
-
-                # Get the camera frame pointmaps
-                curr_view_pts3d_cam = curr_view_prediction["pts3d_local"].unsqueeze(0)
-
-                # Get the depth along ray and ray directions
-                curr_view_depth_along_ray = torch.norm(
-                    curr_view_pts3d_cam, dim=-1, keepdim=True
-                )
-                curr_view_ray_dirs = curr_view_pts3d_cam / curr_view_depth_along_ray
-
-                # Get the pointmaps
-                curr_view_pts3d = curr_view_prediction["pts3d"].unsqueeze(0)
-
-                # Append the outputs to the result list
-                res.append(
-                    {
-                        "pts3d": curr_view_pts3d.to(device),
-                        "pts3d_cam": curr_view_pts3d_cam.to(device),
-                        "ray_directions": curr_view_ray_dirs.to(device),
-                        "depth_along_ray": curr_view_depth_along_ray.to(device),
-                        "cam_trans": curr_view_cam_translations.to(device),
-                        "cam_quats": curr_view_cam_quats.to(device),
-                        "conf": curr_view_conf.to(device),
-                    }
-                )
-
-        return res

mapanything/models/external/pi3/__init__.py

@@ -1,119 +0,0 @@
-"""
-Inference wrapper for Pi3
-"""
-
-import torch
-
-from mapanything.models.external.pi3.models.pi3 import Pi3
-from mapanything.models.external.vggt.utils.rotation import mat_to_quat
-
-
-class Pi3Wrapper(torch.nn.Module):
-    def __init__(
-        self,
-        name,
-        torch_hub_force_reload,
-        load_pretrained_weights=True,
-        pos_type="rope100",
-        decoder_size="large",
-    ):
-        super().__init__()
-        self.name = name
-        self.torch_hub_force_reload = torch_hub_force_reload
-
-        if load_pretrained_weights:
-            # Load pre-trained weights
-            if not torch_hub_force_reload:
-                # Initialize the Pi3 model from huggingface hub cache
-                print("Loading Pi3 from huggingface cache ...")
-                self.model = Pi3.from_pretrained(
-                    "yyfz233/Pi3",
-                )
-            else:
-                # Initialize the Pi3 model
-                self.model = Pi3.from_pretrained("yyfz233/Pi3", force_download=True)
-        else:
-            # Load the Pi3 class
-            self.model = Pi3(
-                pos_type=pos_type,
-                decoder_size=decoder_size,
-            )
-
-        # Get the dtype for Pi3 inference
-        # bfloat16 is supported on Ampere GPUs (Compute Capability 8.0+)
-        self.dtype = (
-            torch.bfloat16
-            if torch.cuda.get_device_capability()[0] >= 8
-            else torch.float16
-        )
-
-    def forward(self, views):
-        """
-        Forward pass wrapper for Pi3
-
-        Assumption:
-        - All the input views have the same image shape.
-
-        Args:
-            views (List[dict]): List of dictionaries containing the input views' images and instance information.
-                                Each dictionary should contain the following keys:
-                                    "img" (tensor): Image tensor of shape (B, C, H, W).
-                                    "data_norm_type" (list): ["identity"]
-
-        Returns:
-            List[dict]: A list containing the final outputs for all N views.
-        """
-        # Get input shape of the images, number of views, and batch size per view
-        batch_size_per_view, _, height, width = views[0]["img"].shape
-        num_views = len(views)
-
-        # Check the data norm type
-        # Pi3 expects a normalized image but without the DINOv2 mean and std applied ("identity")
-        data_norm_type = views[0]["data_norm_type"][0]
-        assert data_norm_type == "identity", (
-            "Pi3 expects a normalized image but without the DINOv2 mean and std applied"
-        )
-
-        # Concatenate the images to create a single (B, V, C, H, W) tensor
-        img_list = [view["img"] for view in views]
-        images = torch.stack(img_list, dim=1)
-
-        # Run the Pi3 aggregator
-        with torch.autocast("cuda", dtype=self.dtype):
-            results = self.model(images)
-
-        # Need high precision for transformations
-        with torch.autocast("cuda", enabled=False):
-            # Convert the output to MapAnything format
-            res = []
-            for view_idx in range(num_views):
-                # Get the extrinsics
-                curr_view_extrinsic = results["camera_poses"][:, view_idx, ...]
-                curr_view_cam_translations = curr_view_extrinsic[..., :3, 3]
-                curr_view_cam_quats = mat_to_quat(curr_view_extrinsic[..., :3, :3])
-
-                # Get the depth along ray, ray directions, local point cloud & global point cloud
-                curr_view_pts3d_cam = results["local_points"][:, view_idx, ...]
-                curr_view_depth_along_ray = torch.norm(
-                    curr_view_pts3d_cam, dim=-1, keepdim=True
-                )
-                curr_view_ray_dirs = curr_view_pts3d_cam / curr_view_depth_along_ray
-                curr_view_pts3d = results["points"][:, view_idx, ...]
-
-                # Get the confidence
-                curr_view_confidence = results["conf"][:, view_idx, ...]
-
-                # Append the outputs to the result list
-                res.append(
-                    {
-                        "pts3d": curr_view_pts3d,
-                        "pts3d_cam": curr_view_pts3d_cam,
-                        "ray_directions": curr_view_ray_dirs,
-                        "depth_along_ray": curr_view_depth_along_ray,
-                        "cam_trans": curr_view_cam_translations,
-                        "cam_quats": curr_view_cam_quats,
-                        "conf": curr_view_confidence,
-                    }
-                )
-
-        return res

mapanything/models/external/pi3/layers/attention.py

@@ -1,429 +0,0 @@
-# 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/master/vision_transformer.py
-#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
-
-
-import os
-
-import torch
-from torch import nn, Tensor
-from torch.nn.attention import SDPBackend
-from torch.nn.functional import scaled_dot_product_attention
-
-XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
-try:
-    if XFORMERS_ENABLED:
-        from xformers.ops import memory_efficient_attention
-
-        XFORMERS_AVAILABLE = True
-        # warnings.warn("xFormers is available (Attention)")
-    else:
-        # warnings.warn("xFormers is disabled (Attention)")
-        raise ImportError
-except ImportError:
-    XFORMERS_AVAILABLE = False
-    # warnings.warn("xFormers is not available (Attention)")
-
-
-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, attn_bias=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]
-        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, attn_bias=None) -> Tensor:
-        if not XFORMERS_AVAILABLE:
-            if attn_bias is not None:
-                raise AssertionError("xFormers is required for using nested tensors")
-            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)
-        q, k, v = [qkv[:, :, i] for i in range(3)]
-
-        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
-        x = x.reshape([B, N, C])
-
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-class FlashAttention(Attention):
-    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
-        B, N, C = x.shape
-        qkv = (
-            self.qkv(x)
-            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
-            .transpose(1, 3)
-        )
-
-        # q, k, v = unbind(qkv, 2)
-        q, k, v = [qkv[:, :, i] for i in range(3)]
-
-        if q.dtype == torch.bfloat16:
-            with nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
-                x = scaled_dot_product_attention(q, k, v)
-        else:
-            with nn.attention.sdpa_kernel(
-                [SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]
-            ):
-                x = scaled_dot_product_attention(q, k, v)
-
-        x = x.transpose(1, 2).reshape([B, N, C])
-
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-"""
-Following is written by GPT-4o
-"""
-
-
-class CrossAttentionRope(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,
-        qk_norm: bool = False,
-        norm_layer: nn.Module = nn.LayerNorm,
-        rope=None,
-    ) -> None:
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        self.scale = head_dim**-0.5
-
-        # Separate projection layers for query, key, and value
-        self.q_proj = nn.Linear(dim, dim, bias=qkv_bias)
-        self.k_proj = nn.Linear(dim, dim, bias=qkv_bias)
-        self.v_proj = nn.Linear(dim, dim, 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,
-        query: Tensor,
-        key: Tensor,
-        value: Tensor,
-        attn_bias=None,
-        qpos=None,
-        kpos=None,
-    ) -> Tensor:
-        """
-        Args:
-            query: Tensor of shape (B, N, C), input query
-            key: Tensor of shape (B, M, C), input key
-            value: Tensor of shape (B, M, C), input value
-            attn_bias: Optional tensor for attention bias
-        Returns:
-            Tensor of shape (B, N, C), output of cross-attention
-        """
-        B, N, C = query.shape
-        _, M, _ = key.shape
-
-        # Project query, key, and value
-        q = (
-            self.q_proj(query)
-            .reshape(B, N, self.num_heads, C // self.num_heads)
-            .permute(0, 2, 1, 3)
-        )
-        k = (
-            self.k_proj(key)
-            .reshape(B, M, self.num_heads, C // self.num_heads)
-            .permute(0, 2, 1, 3)
-        )
-        v = (
-            self.v_proj(value)
-            .reshape(B, M, self.num_heads, C // self.num_heads)
-            .permute(0, 2, 1, 3)
-        )
-        q, k = self.q_norm(q).to(v.dtype), self.k_norm(k).to(v.dtype)
-
-        if self.rope is not None:
-            q = self.rope(q, qpos)
-            k = self.rope(k, kpos)
-
-        # Scale query
-        q = q * self.scale
-
-        # Compute attention scores
-        attn = q @ k.transpose(-2, -1)  # (B, num_heads, N, M)
-        if attn_bias is not None:
-            attn = attn + attn_bias
-
-        attn = attn.softmax(dim=-1)
-        attn = self.attn_drop(attn)
-
-        # Compute attention output
-        x = (attn @ v).transpose(1, 2).reshape(B, N, C)  # (B, N, C)
-
-        # Final projection
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-class MemEffCrossAttentionRope(CrossAttentionRope):
-    def forward(
-        self,
-        query: Tensor,
-        key: Tensor,
-        value: Tensor,
-        attn_bias=None,
-        qpos=None,
-        kpos=None,
-    ) -> Tensor:
-        """
-        Args:
-            query: Tensor of shape (B, N, C), input query
-            key: Tensor of shape (B, M, C), input key
-            value: Tensor of shape (B, M, C), input value
-            attn_bias: Optional tensor for attention bias
-        Returns:
-            Tensor of shape (B, N, C), output of cross-attention
-        """
-        if not XFORMERS_AVAILABLE:
-            if attn_bias is not None:
-                raise AssertionError("xFormers is required for using nested tensors")
-            return super().forward(query, key, value, attn_bias)
-
-        B, N, C = query.shape
-        _, M, _ = key.shape
-
-        # Project query, key, and value
-        q = self.q_proj(query).reshape(B, N, self.num_heads, C // self.num_heads)
-        k = self.k_proj(key).reshape(B, M, self.num_heads, C // self.num_heads)
-        v = self.v_proj(value).reshape(B, M, self.num_heads, C // self.num_heads)
-
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        q, k = self.q_norm(q).to(v.dtype), self.k_norm(k).to(v.dtype)
-
-        if self.rope is not None:
-            q = self.rope(q, qpos)
-            k = self.rope(k, kpos)
-
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-
-        # Compute memory-efficient attention
-        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
-        x = x.reshape(B, N, C)
-
-        # Final projection
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-class AttentionRope(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,
-        qk_norm: bool = False,
-        norm_layer: nn.Module = nn.LayerNorm,
-        rope=None,
-    ) -> 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.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.rope = rope
-
-    def forward(self, x: Tensor, attn_bias=None, xpos=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).to(v.dtype), self.k_norm(k).to(v.dtype)
-
-        if self.rope is not None:
-            q = self.rope(q, xpos)
-            k = self.rope(k, xpos)
-
-        q = q * self.scale
-        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 MemEffAttentionRope(AttentionRope):
-    def forward(self, x: Tensor, attn_bias=None, xpos=None) -> Tensor:
-        if not XFORMERS_AVAILABLE:
-            if attn_bias is not None:
-                raise AssertionError("xFormers is required for using nested tensors")
-            return super().forward(x)
-
-        B, N, C = x.shape
-        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
-
-        qkv = qkv.transpose(1, 3)
-        # q, k, v = unbind(qkv, 2)
-        q, k, v = [qkv[:, :, i] for i in range(3)]
-        q, k = self.q_norm(q).to(v.dtype), self.k_norm(k).to(v.dtype)
-
-        if self.rope is not None:
-            q = self.rope(q, xpos)
-            k = self.rope(k, xpos)
-
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-
-        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
-        x = x.reshape([B, N, C])
-
-        # score_matrix = (q.permute(0, 2, 1, 3) * self.scale @ k.permute(0, 2, 1, 3).transpose(-2, -1)).sum(dim=1).reshape(frame_num, 261, frame_num, 261).mean(dim=[1, 3]).sum(1)         # for frame attention matrix
-        # global_valid_id = torch.where(score_matrix > 0)
-        # score_matrix = (q.permute(0, 2, 1, 3) * self.scale @ k.permute(0, 2, 1, 3).transpose(-2, -1)).sum(dim=1)
-
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-class FlashAttentionRope(AttentionRope):
-    def forward(self, x: Tensor, attn_bias=None, xpos=None) -> Tensor:
-        B, N, C = x.shape
-        qkv = (
-            self.qkv(x)
-            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
-            .transpose(1, 3)
-        )
-
-        # q, k, v = unbind(qkv, 2)
-        q, k, v = [qkv[:, :, i] for i in range(3)]
-        q, k = self.q_norm(q).to(v.dtype), self.k_norm(k).to(v.dtype)
-
-        if self.rope is not None:
-            q = self.rope(q, xpos)
-            k = self.rope(k, xpos)
-
-        if q.dtype == torch.bfloat16:
-            with nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
-                x = scaled_dot_product_attention(q, k, v)
-        else:
-            with nn.attention.sdpa_kernel(
-                [SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]
-            ):
-                x = scaled_dot_product_attention(q, k, v)
-
-        x = x.transpose(1, 2).reshape([B, N, C])
-
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-def get_attn_score(blk_class, x, frame_num, token_length, xpos=None):
-    x = blk_class.norm1(x)
-
-    B, N, C = x.shape
-    qkv = blk_class.attn.qkv(x).reshape(
-        B, N, 3, blk_class.attn.num_heads, C // blk_class.attn.num_heads
-    )
-
-    qkv = qkv.transpose(1, 3)
-    # q, k, v = unbind(qkv, 2)
-    q, k, v = [qkv[:, :, i] for i in range(3)]
-    q, k = blk_class.attn.q_norm(q).to(v.dtype), blk_class.attn.k_norm(k).to(v.dtype)
-
-    if blk_class.attn.rope is not None:
-        q = blk_class.attn.rope(q, xpos)
-        k = blk_class.attn.rope(k, xpos)
-
-    q = q.transpose(1, 2)
-    k = k.transpose(1, 2)
-
-    score = (
-        (
-            q.permute(0, 2, 1, 3)
-            * blk_class.attn.scale
-            @ k.permute(0, 2, 1, 3).transpose(-2, -1)
-        )
-        .sum(dim=1)
-        .reshape(B, frame_num, token_length, frame_num, token_length)
-        .mean(dim=[2, 4])
-        .sum(-1)
-    )
-
-    return score

mapanything/models/external/pi3/layers/block.py

@@ -1,448 +0,0 @@
-# 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/master/vision_transformer.py
-#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
-
-import os
-from typing import Any, Callable, Dict, List, Tuple
-
-import torch
-from torch import nn, Tensor
-
-from mapanything.models.external.dinov2.layers.drop_path import DropPath
-from mapanything.models.external.dinov2.layers.layer_scale import LayerScale
-from mapanything.models.external.dinov2.layers.mlp import Mlp
-from mapanything.models.external.pi3.layers.attention import (
-    Attention,
-    CrossAttentionRope,
-    MemEffAttention,
-)
-
-XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
-try:
-    if XFORMERS_ENABLED:
-        from xformers.ops import fmha, index_select_cat, scaled_index_add
-
-        XFORMERS_AVAILABLE = True
-        # warnings.warn("xFormers is available (Block)")
-    else:
-        # warnings.warn("xFormers is disabled (Block)")
-        raise ImportError
-except ImportError:
-    XFORMERS_AVAILABLE = False
-    # warnings.warn("xFormers is not available (Block)")
-
-
-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,
-    ) -> None:
-        super().__init__()
-        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
-        self.norm1 = norm_layer(dim)
-        self.attn = attn_class(
-            dim,
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            proj_bias=proj_bias,
-            attn_drop=attn_drop,
-            proj_drop=drop,
-        )
-
-        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)
-        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) -> Tensor:
-        def attn_residual_func(x: Tensor) -> Tensor:
-            return self.ls1(self.attn(self.norm1(x)))
-
-        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,
-            )
-            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))
-            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
-        else:
-            x = x + attn_residual_func(x)
-            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,
-) -> 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
-    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
-
-
-def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
-    if scaling_vector is None:
-        x_flat = x.flatten(1)
-        residual = residual.flatten(1)
-        x_plus_residual = torch.index_add(
-            x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor
-        )
-    else:
-        x_plus_residual = scaled_index_add(
-            x,
-            brange,
-            residual.to(dtype=x.dtype),
-            scaling=scaling_vector,
-            alpha=residual_scale_factor,
-        )
-    return x_plus_residual
-
-
-attn_bias_cache: Dict[Tuple, Any] = {}
-
-
-def get_attn_bias_and_cat(x_list, branges=None):
-    """
-    this will perform the index select, cat the tensors, and provide the attn_bias from cache
-    """
-    batch_sizes = (
-        [b.shape[0] for b in branges]
-        if branges is not None
-        else [x.shape[0] for x in x_list]
-    )
-    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
-    if all_shapes not in attn_bias_cache.keys():
-        seqlens = []
-        for b, x in zip(batch_sizes, x_list):
-            for _ in range(b):
-                seqlens.append(x.shape[1])
-        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
-        attn_bias._batch_sizes = batch_sizes
-        attn_bias_cache[all_shapes] = attn_bias
-
-    if branges is not None:
-        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(
-            1, -1, x_list[0].shape[-1]
-        )
-    else:
-        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
-        cat_tensors = torch.cat(tensors_bs1, dim=1)
-
-    return attn_bias_cache[all_shapes], cat_tensors
-
-
-def drop_add_residual_stochastic_depth_list(
-    x_list: List[Tensor],
-    residual_func: Callable[[Tensor, Any], Tensor],
-    sample_drop_ratio: float = 0.0,
-    scaling_vector=None,
-) -> Tensor:
-    # 1) generate random set of indices for dropping samples in the batch
-    branges_scales = [
-        get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list
-    ]
-    branges = [s[0] for s in branges_scales]
-    residual_scale_factors = [s[1] for s in branges_scales]
-
-    # 2) get attention bias and index+concat the tensors
-    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
-
-    # 3) apply residual_func to get residual, and split the result
-    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
-
-    outputs = []
-    for x, brange, residual, residual_scale_factor in zip(
-        x_list, branges, residual_list, residual_scale_factors
-    ):
-        outputs.append(
-            add_residual(
-                x, brange, residual, residual_scale_factor, scaling_vector
-            ).view_as(x)
-        )
-    return outputs
-
-
-class NestedTensorBlock(Block):
-    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
-        """
-        x_list contains a list of tensors to nest together and run
-        """
-        assert isinstance(self.attn, MemEffAttention)
-
-        if self.training and self.sample_drop_ratio > 0.0:
-
-            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.attn(self.norm1(x), attn_bias=attn_bias)
-
-            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.mlp(self.norm2(x))
-
-            x_list = drop_add_residual_stochastic_depth_list(
-                x_list,
-                residual_func=attn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-                scaling_vector=self.ls1.gamma
-                if isinstance(self.ls1, LayerScale)
-                else None,
-            )
-            x_list = drop_add_residual_stochastic_depth_list(
-                x_list,
-                residual_func=ffn_residual_func,
-                sample_drop_ratio=self.sample_drop_ratio,
-                scaling_vector=self.ls2.gamma
-                if isinstance(self.ls1, LayerScale)
-                else None,
-            )
-            return x_list
-        else:
-
-            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
-
-            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
-                return self.ls2(self.mlp(self.norm2(x)))
-
-            attn_bias, x = get_attn_bias_and_cat(x_list)
-            x = x + attn_residual_func(x, attn_bias=attn_bias)
-            x = x + ffn_residual_func(x)
-            return attn_bias.split(x)
-
-    def forward(self, x_or_x_list):
-        if isinstance(x_or_x_list, Tensor):
-            return super().forward(x_or_x_list)
-        elif isinstance(x_or_x_list, list):
-            if not XFORMERS_AVAILABLE:
-                raise AssertionError("xFormers is required for using nested tensors")
-            return self.forward_nested(x_or_x_list)
-        else:
-            raise AssertionError
-
-
-class BlockRope(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,
-    ) -> None:
-        super().__init__()
-        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
-        self.norm1 = norm_layer(dim)
-        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)
-        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, xpos=None) -> Tensor:
-        def attn_residual_func(x: Tensor) -> Tensor:
-            return self.ls1(self.attn(self.norm1(x), xpos=xpos))
-
-        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,
-            )
-            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))
-            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
-        else:
-            x = x + attn_residual_func(x)
-            x = x + ffn_residual_func(x)
-        return x
-
-
-class CrossBlockRope(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,
-        act_layer: Callable[..., nn.Module] = nn.GELU,
-        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
-        attn_class: Callable[..., nn.Module] = Attention,
-        cross_attn_class: Callable[..., nn.Module] = CrossAttentionRope,
-        ffn_layer: Callable[..., nn.Module] = Mlp,
-        init_values=None,
-        qk_norm: bool = False,
-        rope=None,
-    ) -> None:
-        super().__init__()
-        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
-        self.ls1 = (
-            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
-        )
-        self.norm1 = norm_layer(dim)
-        self.attn = attn_class(
-            dim,
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            proj_bias=proj_bias,
-            rope=rope,
-            qk_norm=qk_norm,
-        )
-
-        self.ls2 = (
-            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
-        )
-        self.ls_y = (
-            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
-        )
-        self.norm2 = norm_layer(dim)
-        self.norm_y = norm_layer(dim)
-        self.cross_attn = cross_attn_class(
-            dim,
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            proj_bias=proj_bias,
-            rope=rope,
-            qk_norm=qk_norm,
-        )
-
-        self.norm3 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = ffn_layer(
-            in_features=dim,
-            hidden_features=mlp_hidden_dim,
-            act_layer=act_layer,
-            bias=ffn_bias,
-        )
-
-    def forward(self, x: Tensor, y: Tensor, xpos=None, ypos=None) -> Tensor:
-        def attn_residual_func(x: Tensor) -> Tensor:
-            return self.ls1(self.attn(self.norm1(x), xpos=xpos))
-
-        def cross_attn_residual_func(x: Tensor, y: Tensor) -> Tensor:
-            return self.ls_y(self.cross_attn(self.norm2(x), y, y, qpos=xpos, kpos=ypos))
-
-        def ffn_residual_func(x: Tensor) -> Tensor:
-            return self.ls2(self.mlp(self.norm3(x)))
-
-        x = x + attn_residual_func(x)
-        y_ = self.norm_y(y)
-        x = x + cross_attn_residual_func(x, y_)
-        x = x + ffn_residual_func(x)
-
-        return x

mapanything/models/external/pi3/layers/camera_head.py

@@ -1,106 +0,0 @@
-from copy import deepcopy
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-# code adapted from 'https://github.com/nianticlabs/marepo/blob/9a45e2bb07e5bb8cb997620088d352b439b13e0e/transformer/transformer.py#L172'
-class ResConvBlock(nn.Module):
-    """
-    1x1 convolution residual block
-    """
-
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.head_skip = (
-            nn.Identity()
-            if self.in_channels == self.out_channels
-            else nn.Conv2d(self.in_channels, self.out_channels, 1, 1, 0)
-        )
-        # self.res_conv1 = nn.Conv2d(self.in_channels, self.out_channels, 1, 1, 0)
-        # self.res_conv2 = nn.Conv2d(self.out_channels, self.out_channels, 1, 1, 0)
-        # self.res_conv3 = nn.Conv2d(self.out_channels, self.out_channels, 1, 1, 0)
-
-        # change 1x1 convolution to linear
-        self.res_conv1 = nn.Linear(self.in_channels, self.out_channels)
-        self.res_conv2 = nn.Linear(self.out_channels, self.out_channels)
-        self.res_conv3 = nn.Linear(self.out_channels, self.out_channels)
-
-    def forward(self, res):
-        x = F.relu(self.res_conv1(res))
-        x = F.relu(self.res_conv2(x))
-        x = F.relu(self.res_conv3(x))
-        res = self.head_skip(res) + x
-        return res
-
-
-class CameraHead(nn.Module):
-    def __init__(self, dim=512):
-        super().__init__()
-        output_dim = dim
-        self.res_conv = nn.ModuleList(
-            [deepcopy(ResConvBlock(output_dim, output_dim)) for _ in range(2)]
-        )
-        self.avgpool = nn.AdaptiveAvgPool2d(1)
-        self.more_mlps = nn.Sequential(
-            nn.Linear(output_dim, output_dim),
-            nn.ReLU(),
-            nn.Linear(output_dim, output_dim),
-            nn.ReLU(),
-        )
-        self.fc_t = nn.Linear(output_dim, 3)
-        self.fc_rot = nn.Linear(output_dim, 9)
-
-    def forward(self, feat, patch_h, patch_w):
-        BN, hw, c = feat.shape
-
-        for i in range(2):
-            feat = self.res_conv[i](feat)
-
-        # feat = self.avgpool(feat)
-        feat = self.avgpool(
-            feat.permute(0, 2, 1).reshape(BN, -1, patch_h, patch_w).contiguous()
-        )  ##########
-        feat = feat.view(feat.size(0), -1)
-
-        feat = self.more_mlps(feat)  # [B, D_]
-        with torch.amp.autocast(device_type="cuda", enabled=False):
-            out_t = self.fc_t(feat.float())  # [B,3]
-            out_r = self.fc_rot(feat.float())  # [B,9]
-            pose = self.convert_pose_to_4x4(BN, out_r, out_t, feat.device)
-
-        return pose
-
-    def convert_pose_to_4x4(self, B, out_r, out_t, device):
-        out_r = self.svd_orthogonalize(out_r)  # [N,3,3]
-        pose = torch.zeros((B, 4, 4), device=device)
-        pose[:, :3, :3] = out_r
-        pose[:, :3, 3] = out_t
-        pose[:, 3, 3] = 1.0
-        return pose
-
-    def svd_orthogonalize(self, m):
-        """Convert 9D representation to SO(3) using SVD orthogonalization.
-
-        Args:
-          m: [BATCH, 3, 3] 3x3 matrices.
-
-        Returns:
-          [BATCH, 3, 3] SO(3) rotation matrices.
-        """
-        if m.dim() < 3:
-            m = m.reshape((-1, 3, 3))
-        m_transpose = torch.transpose(
-            torch.nn.functional.normalize(m, p=2, dim=-1), dim0=-1, dim1=-2
-        )
-        u, s, v = torch.svd(m_transpose)
-        det = torch.det(torch.matmul(v, u.transpose(-2, -1)))
-        # Check orientation reflection.
-        r = torch.matmul(
-            torch.cat([v[:, :, :-1], v[:, :, -1:] * det.view(-1, 1, 1)], dim=2),
-            u.transpose(-2, -1),
-        )
-        return r

mapanything/models/external/pi3/layers/pos_embed.py

@@ -1,190 +0,0 @@
-# Copyright (C) 2022-present Naver Corporation. All rights reserved.
-# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
-
-
-# --------------------------------------------------------
-# Position embedding utils
-# --------------------------------------------------------
-
-
-import numpy as np
-import torch
-
-
-# --------------------------------------------------------
-# 2D sine-cosine position embedding
-# References:
-# MAE: https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
-# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
-# MoCo v3: https://github.com/facebookresearch/moco-v3
-# --------------------------------------------------------
-def get_2d_sincos_pos_embed(embed_dim, grid_size, n_cls_token=0):
-    """
-    grid_size: int of the grid height and width
-    return:
-    pos_embed: [grid_size*grid_size, embed_dim] or [n_cls_token+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
-    """
-    grid_h = np.arange(grid_size, dtype=np.float32)
-    grid_w = np.arange(grid_size, dtype=np.float32)
-    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
-    grid = np.stack(grid, axis=0)
-
-    grid = grid.reshape([2, 1, grid_size, grid_size])
-    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
-    if n_cls_token > 0:
-        pos_embed = np.concatenate(
-            [np.zeros([n_cls_token, embed_dim]), pos_embed], axis=0
-        )
-    return pos_embed
-
-
-def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
-    assert embed_dim % 2 == 0
-
-    # use half of dimensions to encode grid_h
-    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
-    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
-
-    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
-    return emb
-
-
-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
-
-
-# --------------------------------------------------------
-# Interpolate position embeddings for high-resolution
-# References:
-# MAE: https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
-# DeiT: https://github.com/facebookresearch/deit
-# --------------------------------------------------------
-def interpolate_pos_embed(model, checkpoint_model):
-    if "pos_embed" in checkpoint_model:
-        pos_embed_checkpoint = checkpoint_model["pos_embed"]
-        embedding_size = pos_embed_checkpoint.shape[-1]
-        num_patches = model.patch_embed.num_patches
-        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
-        # height (== width) for the checkpoint position embedding
-        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
-        # height (== width) for the new position embedding
-        new_size = int(num_patches**0.5)
-        # class_token and dist_token are kept unchanged
-        if orig_size != new_size:
-            print(
-                "Position interpolate from %dx%d to %dx%d"
-                % (orig_size, orig_size, new_size, new_size)
-            )
-            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
-            # only the position tokens are interpolated
-            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
-            pos_tokens = pos_tokens.reshape(
-                -1, orig_size, orig_size, embedding_size
-            ).permute(0, 3, 1, 2)
-            pos_tokens = torch.nn.functional.interpolate(
-                pos_tokens,
-                size=(new_size, new_size),
-                mode="bicubic",
-                align_corners=False,
-            )
-            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
-            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
-            checkpoint_model["pos_embed"] = new_pos_embed
-
-
-# ----------------------------------------------------------
-# RoPE2D: RoPE implementation in 2D
-# ----------------------------------------------------------
-
-try:
-    from models.curope import cuRoPE2D
-
-    RoPE2D = cuRoPE2D
-except ImportError:
-
-    class RoPE2D(torch.nn.Module):
-        def __init__(self, freq=100.0, F0=1.0):
-            super().__init__()
-            self.base = freq
-            self.F0 = F0
-            self.cache = {}
-
-        def get_cos_sin(self, D, seq_len, device, dtype):
-            if (D, seq_len, device, dtype) not in self.cache:
-                inv_freq = 1.0 / (
-                    self.base ** (torch.arange(0, D, 2).float().to(device) / D)
-                )
-                t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
-                freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
-                freqs = torch.cat((freqs, freqs), dim=-1)
-                cos = freqs.cos()  # (Seq, Dim)
-                sin = freqs.sin()
-                self.cache[D, seq_len, device, dtype] = (cos, sin)
-            return self.cache[D, seq_len, device, dtype]
-
-        @staticmethod
-        def rotate_half(x):
-            x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
-            return torch.cat((-x2, x1), dim=-1)
-
-        def apply_rope1d(self, tokens, pos1d, cos, sin):
-            assert pos1d.ndim == 2
-            cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :]
-            sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :]
-            return (tokens * cos) + (self.rotate_half(tokens) * sin)
-
-        def forward(self, tokens, positions):
-            """
-            input:
-                * tokens: batch_size x nheads x ntokens x dim
-                * positions: batch_size x ntokens x 2 (y and x position of each token)
-            output:
-                * tokens after appplying RoPE2D (batch_size x nheads x ntokens x dim)
-            """
-            assert tokens.size(3) % 2 == 0, (
-                "number of dimensions should be a multiple of two"
-            )
-            D = tokens.size(3) // 2
-            assert positions.ndim == 3 and positions.shape[-1] == 2  # Batch, Seq, 2
-            cos, sin = self.get_cos_sin(
-                D, int(positions.max()) + 1, tokens.device, tokens.dtype
-            )
-            # split features into two along the feature dimension, and apply rope1d on each half
-            y, x = tokens.chunk(2, dim=-1)
-            y = self.apply_rope1d(y, positions[:, :, 0], cos, sin)
-            x = self.apply_rope1d(x, positions[:, :, 1], cos, sin)
-            tokens = torch.cat((y, x), dim=-1)
-            return tokens
-
-
-# patch embedding
-class PositionGetter(object):
-    """return positions of patches"""
-
-    def __init__(self):
-        self.cache_positions = {}
-
-    def __call__(self, b, h, w, device):
-        if (h, w) not in self.cache_positions:
-            x = torch.arange(w, device=device)
-            y = torch.arange(h, device=device)
-            self.cache_positions[h, w] = torch.cartesian_prod(y, x)  # (h, w, 2)
-        pos = self.cache_positions[h, w].view(1, h * w, 2).expand(b, -1, 2).clone()
-        return pos

mapanything/models/external/pi3/layers/transformer_head.py

@@ -1,98 +0,0 @@
-from functools import partial
-
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.utils.checkpoint import checkpoint
-
-from mapanything.models.external.dinov2.layers import Mlp
-from mapanything.models.external.pi3.layers.attention import FlashAttentionRope
-from mapanything.models.external.pi3.layers.block import BlockRope
-
-
-class TransformerDecoder(nn.Module):
-    def __init__(
-        self,
-        in_dim,
-        out_dim,
-        dec_embed_dim=512,
-        depth=5,
-        dec_num_heads=8,
-        mlp_ratio=4,
-        rope=None,
-        need_project=True,
-        use_checkpoint=False,
-    ):
-        super().__init__()
-
-        self.projects = (
-            nn.Linear(in_dim, dec_embed_dim) if need_project else nn.Identity()
-        )
-        self.use_checkpoint = use_checkpoint
-
-        self.blocks = nn.ModuleList(
-            [
-                BlockRope(
-                    dim=dec_embed_dim,
-                    num_heads=dec_num_heads,
-                    mlp_ratio=mlp_ratio,
-                    qkv_bias=True,
-                    proj_bias=True,
-                    ffn_bias=True,
-                    drop_path=0.0,
-                    norm_layer=partial(nn.LayerNorm, eps=1e-6),
-                    act_layer=nn.GELU,
-                    ffn_layer=Mlp,
-                    init_values=None,
-                    qk_norm=False,
-                    # attn_class=MemEffAttentionRope,
-                    attn_class=FlashAttentionRope,
-                    rope=rope,
-                )
-                for _ in range(depth)
-            ]
-        )
-
-        self.linear_out = nn.Linear(dec_embed_dim, out_dim)
-
-    def forward(self, hidden, xpos=None):
-        hidden = self.projects(hidden)
-        for i, blk in enumerate(self.blocks):
-            if self.use_checkpoint and self.training:
-                hidden = checkpoint(blk, hidden, xpos=xpos, use_reentrant=False)
-            else:
-                hidden = blk(hidden, xpos=xpos)
-        out = self.linear_out(hidden)
-        return out
-
-
-class LinearPts3d(nn.Module):
-    """
-    Linear head for dust3r
-    Each token outputs: - 16x16 3D points (+ confidence)
-    """
-
-    def __init__(
-        self,
-        patch_size,
-        dec_embed_dim,
-        output_dim=3,
-    ):
-        super().__init__()
-        self.patch_size = patch_size
-
-        self.proj = nn.Linear(dec_embed_dim, (output_dim) * self.patch_size**2)
-
-    def forward(self, decout, img_shape):
-        H, W = img_shape
-        tokens = decout[-1]
-        B, S, D = tokens.shape
-
-        # extract 3D points
-        feat = self.proj(tokens)  # B,S,D
-        feat = feat.transpose(-1, -2).view(
-            B, -1, H // self.patch_size, W // self.patch_size
-        )
-        feat = F.pixel_shuffle(feat, self.patch_size)  # B,3,H,W
-
-        # permute + norm depth
-        return feat.permute(0, 2, 3, 1)

mapanything/models/external/pi3/models/pi3.py

@@ -1,251 +0,0 @@
-from copy import deepcopy
-from functools import partial
-
-import torch
-import torch.nn as nn
-from huggingface_hub import PyTorchModelHubMixin
-
-from mapanything.models.external.dinov2.hub.backbones import dinov2_vitl14_reg
-from mapanything.models.external.dinov2.layers import Mlp
-from mapanything.models.external.pi3.layers.attention import FlashAttentionRope
-from mapanything.models.external.pi3.layers.block import BlockRope
-from mapanything.models.external.pi3.layers.camera_head import CameraHead
-from mapanything.models.external.pi3.layers.pos_embed import PositionGetter, RoPE2D
-from mapanything.models.external.pi3.layers.transformer_head import (
-    LinearPts3d,
-    TransformerDecoder,
-)
-
-
-def homogenize_points(
-    points,
-):
-    """Convert batched points (xyz) to (xyz1)."""
-    return torch.cat([points, torch.ones_like(points[..., :1])], dim=-1)
-
-
-class Pi3(nn.Module, PyTorchModelHubMixin):
-    def __init__(
-        self,
-        pos_type="rope100",
-        decoder_size="large",
-    ):
-        super().__init__()
-
-        # ----------------------
-        #        Encoder
-        # ----------------------
-        self.encoder = dinov2_vitl14_reg(pretrained=False)
-        self.patch_size = 14
-        del self.encoder.mask_token
-
-        # ----------------------
-        #  Positonal Encoding
-        # ----------------------
-        self.pos_type = pos_type if pos_type is not None else "none"
-        self.rope = None
-        if self.pos_type.startswith("rope"):  # eg rope100
-            if RoPE2D is None:
-                raise ImportError(
-                    "Cannot find cuRoPE2D, please install it following the README instructions"
-                )
-            freq = float(self.pos_type[len("rope") :])
-            self.rope = RoPE2D(freq=freq)
-            self.position_getter = PositionGetter()
-        else:
-            raise NotImplementedError
-
-        # ----------------------
-        #        Decoder
-        # ----------------------
-        if decoder_size == "small":
-            dec_embed_dim = 384
-            dec_num_heads = 6
-            mlp_ratio = 4
-            dec_depth = 24
-        elif decoder_size == "base":
-            dec_embed_dim = 768
-            dec_num_heads = 12
-            mlp_ratio = 4
-            dec_depth = 24
-        elif decoder_size == "large":
-            dec_embed_dim = 1024
-            dec_num_heads = 16
-            mlp_ratio = 4
-            dec_depth = 36
-        else:
-            raise NotImplementedError
-        self.decoder = nn.ModuleList(
-            [
-                BlockRope(
-                    dim=dec_embed_dim,
-                    num_heads=dec_num_heads,
-                    mlp_ratio=mlp_ratio,
-                    qkv_bias=True,
-                    proj_bias=True,
-                    ffn_bias=True,
-                    drop_path=0.0,
-                    norm_layer=partial(nn.LayerNorm, eps=1e-6),
-                    act_layer=nn.GELU,
-                    ffn_layer=Mlp,
-                    init_values=0.01,
-                    qk_norm=True,
-                    attn_class=FlashAttentionRope,
-                    rope=self.rope,
-                )
-                for _ in range(dec_depth)
-            ]
-        )
-        self.dec_embed_dim = dec_embed_dim
-
-        # ----------------------
-        #     Register_token
-        # ----------------------
-        num_register_tokens = 5
-        self.patch_start_idx = num_register_tokens
-        self.register_token = nn.Parameter(
-            torch.randn(1, 1, num_register_tokens, self.dec_embed_dim)
-        )
-        nn.init.normal_(self.register_token, std=1e-6)
-
-        # ----------------------
-        #  Local Points Decoder
-        # ----------------------
-        self.point_decoder = TransformerDecoder(
-            in_dim=2 * self.dec_embed_dim,
-            dec_embed_dim=1024,
-            dec_num_heads=16,
-            out_dim=1024,
-            rope=self.rope,
-        )
-        self.point_head = LinearPts3d(patch_size=14, dec_embed_dim=1024, output_dim=3)
-
-        # ----------------------
-        #     Conf Decoder
-        # ----------------------
-        self.conf_decoder = deepcopy(self.point_decoder)
-        self.conf_head = LinearPts3d(patch_size=14, dec_embed_dim=1024, output_dim=1)
-
-        # ----------------------
-        #  Camera Pose Decoder
-        # ----------------------
-        self.camera_decoder = TransformerDecoder(
-            in_dim=2 * self.dec_embed_dim,
-            dec_embed_dim=1024,
-            dec_num_heads=16,  # 8
-            out_dim=512,
-            rope=self.rope,
-            use_checkpoint=False,
-        )
-        self.camera_head = CameraHead(dim=512)
-
-        # For ImageNet Normalize
-        image_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
-        image_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
-
-        self.register_buffer("image_mean", image_mean)
-        self.register_buffer("image_std", image_std)
-
-    def decode(self, hidden, N, H, W):
-        BN, hw, _ = hidden.shape
-        B = BN // N
-
-        final_output = []
-
-        hidden = hidden.reshape(B * N, hw, -1)
-
-        register_token = self.register_token.repeat(B, N, 1, 1).reshape(
-            B * N, *self.register_token.shape[-2:]
-        )
-
-        # Concatenate special tokens with patch tokens
-        hidden = torch.cat([register_token, hidden], dim=1)
-        hw = hidden.shape[1]
-
-        if self.pos_type.startswith("rope"):
-            pos = self.position_getter(
-                B * N, H // self.patch_size, W // self.patch_size, hidden.device
-            )
-
-        if self.patch_start_idx > 0:
-            # do not use position embedding for special tokens (camera and register tokens)
-            # so set pos to 0 for the special tokens
-            pos = pos + 1
-            pos_special = (
-                torch.zeros(B * N, self.patch_start_idx, 2)
-                .to(hidden.device)
-                .to(pos.dtype)
-            )
-            pos = torch.cat([pos_special, pos], dim=1)
-
-        for i in range(len(self.decoder)):
-            blk = self.decoder[i]
-
-            if i % 2 == 0:
-                pos = pos.reshape(B * N, hw, -1)
-                hidden = hidden.reshape(B * N, hw, -1)
-            else:
-                pos = pos.reshape(B, N * hw, -1)
-                hidden = hidden.reshape(B, N * hw, -1)
-
-            hidden = blk(hidden, xpos=pos)
-
-            if i + 1 in [len(self.decoder) - 1, len(self.decoder)]:
-                final_output.append(hidden.reshape(B * N, hw, -1))
-
-        return torch.cat([final_output[0], final_output[1]], dim=-1), pos.reshape(
-            B * N, hw, -1
-        )
-
-    def forward(self, imgs):
-        imgs = (imgs - self.image_mean) / self.image_std
-
-        B, N, _, H, W = imgs.shape
-        patch_h, patch_w = H // 14, W // 14
-
-        # encode by dinov2
-        imgs = imgs.reshape(B * N, _, H, W)
-        hidden = self.encoder(imgs, is_training=True)
-
-        if isinstance(hidden, dict):
-            hidden = hidden["x_norm_patchtokens"]
-
-        hidden, pos = self.decode(hidden, N, H, W)
-
-        point_hidden = self.point_decoder(hidden, xpos=pos)
-        conf_hidden = self.conf_decoder(hidden, xpos=pos)
-        camera_hidden = self.camera_decoder(hidden, xpos=pos)
-
-        with torch.amp.autocast(device_type="cuda", enabled=False):
-            # local points
-            point_hidden = point_hidden.float()
-            ret = self.point_head(
-                [point_hidden[:, self.patch_start_idx :]], (H, W)
-            ).reshape(B, N, H, W, -1)
-            xy, z = ret.split([2, 1], dim=-1)
-            z = torch.exp(z)
-            local_points = torch.cat([xy * z, z], dim=-1)
-
-            # confidence
-            conf_hidden = conf_hidden.float()
-            conf = self.conf_head(
-                [conf_hidden[:, self.patch_start_idx :]], (H, W)
-            ).reshape(B, N, H, W, -1)
-
-            # camera
-            camera_hidden = camera_hidden.float()
-            camera_poses = self.camera_head(
-                camera_hidden[:, self.patch_start_idx :], patch_h, patch_w
-            ).reshape(B, N, 4, 4)
-
-            # unproject local points using camera poses
-            points = torch.einsum(
-                "bnij, bnhwj -> bnhwi", camera_poses, homogenize_points(local_points)
-            )[..., :3]
-
-        return dict(
-            points=points,
-            local_points=local_points,
-            conf=conf,
-            camera_poses=camera_poses,
-        )

mapanything/models/external/pow3r/__init__.py

@@ -1,860 +0,0 @@
-"""
-Inference wrapper for Pow3R
-"""
-
-import warnings
-from copy import deepcopy
-
-import pow3r.model.blocks  # noqa
-import roma
-import torch
-import torch.nn as nn
-import tqdm
-from dust3r.cloud_opt import global_aligner, GlobalAlignerMode
-from dust3r.image_pairs import make_pairs
-from dust3r.inference import check_if_same_size
-from dust3r.model import CroCoNet
-from dust3r.patch_embed import get_patch_embed as dust3r_patch_embed
-from dust3r.utils.device import collate_with_cat, to_cpu
-from dust3r.utils.misc import (
-    fill_default_args,
-    freeze_all_params,
-    interleave,
-    is_symmetrized,
-    transpose_to_landscape,
-)
-from pow3r.model.blocks import Block, BlockInject, DecoderBlock, DecoderBlockInject, Mlp
-from pow3r.model.heads import head_factory
-from pow3r.model.inference import (
-    add_depth,
-    add_intrinsics,
-    add_relpose,
-)
-from pow3r.model.patch_embed import get_patch_embed
-
-from mapanything.models.external.vggt.utils.rotation import mat_to_quat
-from mapanything.utils.geometry import (
-    convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap,
-    convert_z_depth_to_depth_along_ray,
-    depthmap_to_camera_frame,
-    get_rays_in_camera_frame,
-)
-
-
-class Pow3R(CroCoNet):
-    """Two siamese encoders, followed by two decoders.
-    The goal is to output 3d points directly, both images in view1's frame
-    (hence the asymmetry).
-    """
-
-    def __init__(
-        self,
-        mode="embed",
-        head_type="linear",
-        patch_embed_cls="PatchEmbedDust3R",
-        freeze="none",
-        landscape_only=True,
-        **croco_kwargs,
-    ):
-        # retrieve all default arguments using python magic
-        self.croco_args = fill_default_args(croco_kwargs, super().__init__)
-        super().__init__(**croco_kwargs)
-        del self.mask_token  # useless
-        del self.prediction_head
-
-        dec_dim, enc_dim = self.decoder_embed.weight.shape
-        self.enc_embed_dim = enc_dim
-        self.dec_embed_dim = dec_dim
-
-        self.mode = mode
-        # additional parameters in the encoder
-        img_size = self.patch_embed.img_size
-        patch_size = self.patch_embed.patch_size[0]
-        self.patch_embed = dust3r_patch_embed(
-            patch_embed_cls, img_size, patch_size, self.enc_embed_dim
-        )
-        self.patch_embed_rays = get_patch_embed(
-            patch_embed_cls + "_Mlp",
-            img_size,
-            patch_size,
-            self.enc_embed_dim,
-            in_chans=3,
-        )
-        self.patch_embed_depth = get_patch_embed(
-            patch_embed_cls + "_Mlp",
-            img_size,
-            patch_size,
-            self.enc_embed_dim,
-            in_chans=2,
-        )
-        self.pose_embed = Mlp(12, 4 * dec_dim, dec_dim)
-
-        # additional parameters in the decoder
-        self.dec_cls = "_cls" in self.mode
-        self.dec_num_cls = 0
-        if self.dec_cls:
-            # use a CLS token in the decoder only
-            self.mode = self.mode.replace("_cls", "")
-            self.cls_token1 = nn.Parameter(torch.zeros((dec_dim,)))
-            self.cls_token2 = nn.Parameter(torch.zeros((dec_dim,)))
-            self.dec_num_cls = 1  # affects all blocks
-
-        use_ln = "_ln" in self.mode  # TODO remove?
-        self.patch_ln = nn.LayerNorm(enc_dim) if use_ln else nn.Identity()
-        self.dec1_pre_ln = nn.LayerNorm(dec_dim) if use_ln else nn.Identity()
-        self.dec2_pre_ln = nn.LayerNorm(dec_dim) if use_ln else nn.Identity()
-
-        self.dec_blocks2 = deepcopy(self.dec_blocks)
-
-        # here we modify some of the blocks
-        self.replace_some_blocks()
-
-        self.set_downstream_head(head_type, landscape_only, **croco_kwargs)
-        self.set_freeze(freeze)
-
-    def replace_some_blocks(self):
-        assert self.mode.startswith("inject")  # inject[0,0.5]
-        NewBlock = BlockInject
-        DecoderNewBlock = DecoderBlockInject
-
-        all_layers = {
-            i / n
-            for i in range(len(self.enc_blocks))
-            for n in [len(self.enc_blocks), len(self.dec_blocks)]
-        }
-        which_layers = eval(self.mode[self.mode.find("[") :]) or all_layers
-        assert isinstance(which_layers, (set, list))
-
-        n = 0
-        for i, block in enumerate(self.enc_blocks):
-            if i / len(self.enc_blocks) in which_layers:
-                block.__class__ = NewBlock
-                block.init(self.enc_embed_dim)
-                n += 1
-            else:
-                block.__class__ = Block
-        assert n == len(which_layers), breakpoint()
-
-        n = 0
-        for i in range(len(self.dec_blocks)):
-            for blocks in [self.dec_blocks, self.dec_blocks2]:
-                block = blocks[i]
-                if i / len(self.dec_blocks) in which_layers:
-                    block.__class__ = DecoderNewBlock
-                    block.init(self.dec_embed_dim)
-                    n += 1
-                else:
-                    block.__class__ = DecoderBlock
-        assert n == 2 * len(which_layers), breakpoint()
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_path, **kw):
-        return _load_model(pretrained_model_path, device="cpu")
-
-    def load_state_dict(self, ckpt, **kw):
-        # duplicate all weights for the second decoder if not present
-        new_ckpt = dict(ckpt)
-        if not any(k.startswith("dec_blocks2") for k in ckpt):
-            for key, value in ckpt.items():
-                if key.startswith("dec_blocks"):
-                    new_ckpt[key.replace("dec_blocks", "dec_blocks2")] = value
-        # remove layers that have different shapes
-        cur_ckpt = self.state_dict()
-        for key, val in ckpt.items():
-            if key.startswith("downstream_head2.proj"):
-                if key in cur_ckpt and cur_ckpt[key].shape != val.shape:
-                    print(f" (removing ckpt[{key}] because wrong shape)")
-                    del new_ckpt[key]
-        return super().load_state_dict(new_ckpt, **kw)
-
-    def set_freeze(self, freeze):  # this is for use by downstream models
-        self.freeze = freeze
-        to_be_frozen = {
-            "none": [],
-            "encoder": [self.patch_embed, self.enc_blocks],
-        }
-        freeze_all_params(to_be_frozen[freeze])
-
-    def set_prediction_head(self, *args, **kwargs):
-        """No prediction head"""
-        return
-
-    def set_downstream_head(
-        self,
-        head_type,
-        landscape_only,
-        patch_size,
-        img_size,
-        mlp_ratio,
-        dec_depth,
-        **kw,
-    ):
-        assert img_size[0] % patch_size == 0 and img_size[1] % patch_size == 0, (
-            f"{img_size=} must be multiple of {patch_size=}"
-        )
-
-        # split heads if different
-        heads = head_type.split(";")
-        assert len(heads) in (1, 2)
-        head1_type, head2_type = (heads + heads)[:2]
-
-        # allocate heads
-        self.downstream_head1 = head_factory(head1_type, self)
-        self.downstream_head2 = head_factory(head2_type, self)
-
-        # magic wrapper
-        self.head1 = transpose_to_landscape(
-            self.downstream_head1, activate=landscape_only
-        )
-        self.head2 = transpose_to_landscape(
-            self.downstream_head2, activate=landscape_only
-        )
-
-    def _encode_image(self, image, true_shape, rays=None, depth=None):
-        # embed the image into patches  (x has size B x Npatches x C)
-        x, pos = self.patch_embed(image, true_shape=true_shape)
-
-        if rays is not None:  # B,3,H,W
-            rays_emb, pos2 = self.patch_embed_rays(rays, true_shape=true_shape)
-            assert (pos == pos2).all()
-            if self.mode.startswith("embed"):
-                x = x + rays_emb
-        else:
-            rays_emb = None
-
-        if depth is not None:  # B,2,H,W
-            depth_emb, pos2 = self.patch_embed_depth(depth, true_shape=true_shape)
-            assert (pos == pos2).all()
-            if self.mode.startswith("embed"):
-                x = x + depth_emb
-        else:
-            depth_emb = None
-
-        x = self.patch_ln(x)
-
-        # add positional embedding without cls token
-        assert self.enc_pos_embed is None
-
-        # now apply the transformer encoder and normalization
-        for blk in self.enc_blocks:
-            x = blk(x, pos, rays=rays_emb, depth=depth_emb)
-
-        x = self.enc_norm(x)
-        return x, pos
-
-    def encode_symmetrized(self, view1, view2):
-        img1 = view1["img"]
-        img2 = view2["img"]
-        B = img1.shape[0]
-        # Recover true_shape when available, otherwise assume that the img shape is the true one
-        shape1 = view1.get(
-            "true_shape", torch.tensor(img1.shape[-2:])[None].repeat(B, 1)
-        )
-        shape2 = view2.get(
-            "true_shape", torch.tensor(img2.shape[-2:])[None].repeat(B, 1)
-        )
-        # warning! maybe the images have different portrait/landscape orientations
-
-        # privileged information
-        rays1 = view1.get("known_rays", None)
-        rays2 = view2.get("known_rays", None)
-        depth1 = view1.get("known_depth", None)
-        depth2 = view2.get("known_depth", None)
-
-        if is_symmetrized(view1, view2):
-            # computing half of forward pass!'
-            def hsub(x):
-                return None if x is None else x[::2]
-
-            feat1, pos1 = self._encode_image(
-                img1[::2], shape1[::2], rays=hsub(rays1), depth=hsub(depth1)
-            )
-            feat2, pos2 = self._encode_image(
-                img2[::2], shape2[::2], rays=hsub(rays2), depth=hsub(depth2)
-            )
-
-            feat1, feat2 = interleave(feat1, feat2)
-            pos1, pos2 = interleave(pos1, pos2)
-        else:
-            feat1, pos1 = self._encode_image(img1, shape1, rays=rays1, depth=depth1)
-            feat2, pos2 = self._encode_image(img2, shape2, rays=rays2, depth=depth2)
-
-        return (shape1, shape2), (feat1, feat2), (pos1, pos2)
-
-    def _decoder(self, f1, pos1, f2, pos2, relpose1=None, relpose2=None):
-        final_output = [(f1, f2)]  # before projection
-
-        # project to decoder dim
-        f1 = self.decoder_embed(f1)
-        f2 = self.decoder_embed(f2)
-
-        # add CLS token for the decoder
-        if self.dec_cls:
-            cls1 = self.cls_token1[None, None].expand(len(f1), 1, -1).clone()
-            cls2 = self.cls_token2[None, None].expand(len(f2), 1, -1).clone()
-
-        if relpose1 is not None:  # shape = (B, 4, 4)
-            pose_emb1 = self.pose_embed(relpose1[:, :3].flatten(1)).unsqueeze(1)
-            if self.mode.startswith("embed"):
-                if self.dec_cls:
-                    cls1 = cls1 + pose_emb1
-                else:
-                    f1 = f1 + pose_emb1
-        else:
-            pose_emb1 = None
-
-        if relpose2 is not None:  # shape = (B, 4, 4)
-            pose_emb2 = self.pose_embed(relpose2[:, :3].flatten(1)).unsqueeze(1)
-            if self.mode.startswith("embed"):
-                if self.dec_cls:
-                    cls2 = cls2 + pose_emb2
-                else:
-                    f2 = f2 + pose_emb2
-        else:
-            pose_emb2 = None
-
-        if self.dec_cls:
-            f1, pos1 = cat_cls(cls1, f1, pos1)
-            f2, pos2 = cat_cls(cls2, f2, pos2)
-
-        f1 = self.dec1_pre_ln(f1)
-        f2 = self.dec2_pre_ln(f2)
-
-        final_output.append((f1, f2))  # to be removed later
-        for blk1, blk2 in zip(self.dec_blocks, self.dec_blocks2):
-            # img1 side
-            f1, _ = blk1(
-                *final_output[-1][::+1],
-                pos1,
-                pos2,
-                relpose=pose_emb1,
-                num_cls=self.dec_num_cls,
-            )
-            # img2 side
-            f2, _ = blk2(
-                *final_output[-1][::-1],
-                pos2,
-                pos1,
-                relpose=pose_emb2,
-                num_cls=self.dec_num_cls,
-            )
-            # store the result
-            final_output.append((f1, f2))
-
-        del final_output[1]  # duplicate with final_output[0] (after decoder proj)
-        if self.dec_cls:  # remove cls token for decoder layers
-            final_output[1:] = [(f1[:, 1:], f2[:, 1:]) for f1, f2 in final_output[1:]]
-        # normalize last output
-        final_output[-1] = tuple(map(self.dec_norm, final_output[-1]))
-        return zip(*final_output)
-
-    def _downstream_head(self, head_num, decout, img_shape):
-        B, S, D = decout[-1].shape
-        head = getattr(self, f"head{head_num}")
-        return head(decout, img_shape)
-
-    def forward(self, view1, view2):
-        # encode the two images --> B,S,D
-        (shape1, shape2), (feat1, feat2), (pos1, pos2) = self.encode_symmetrized(
-            view1, view2
-        )
-
-        # combine all ref images into object-centric representation
-        dec1, dec2 = self._decoder(
-            feat1,
-            pos1,
-            feat2,
-            pos2,
-            relpose1=view1.get("known_pose"),
-            relpose2=view2.get("known_pose"),
-        )
-        with torch.autocast("cuda", enabled=False):
-            res1 = self._downstream_head(1, [tok.float() for tok in dec1], shape1)
-            res2 = self._downstream_head(2, [tok.float() for tok in dec2], shape2)
-
-        res2["pts3d_in_other_view"] = res2.pop(
-            "pts3d"
-        )  # predict view2's pts3d in view1's frame
-        return res1, res2
-
-
-def convert_release_dust3r_args(args):
-    args.model = (
-        args.model.replace("patch_embed_cls", "patch_embed")
-        .replace("AsymmetricMASt3R", "AsymmetricCroCo3DStereo")
-        .replace("PatchEmbedDust3R", "convManyAR")
-        .replace(
-            "pos_embed='RoPE100'",
-            "enc_pos_embed='cuRoPE100', dec_pos_embed='cuRoPE100'",
-        )
-    )
-    return args
-
-
-def _load_model(model_path, device):
-    print("... loading model from", model_path)
-    ckpt = torch.load(model_path, map_location="cpu")
-    try:
-        net = eval(
-            ckpt["args"].model[:-1].replace("convManyAR", "convP")
-            + ", landscape_only=False)"
-        )
-    except Exception:
-        args = convert_release_dust3r_args(ckpt["args"])
-        net = eval(
-            args.model[:-1].replace("convManyAR", "convP") + ", landscape_only=False)"
-        )
-    ckpt["model"] = {
-        k.replace("_downstream_head", "downstream_head"): v
-        for k, v in ckpt["model"].items()
-    }
-    print(net.load_state_dict(ckpt["model"], strict=False))
-    return net.to(device)
-
-
-def cat_cls(cls, tokens, pos):
-    tokens = torch.cat((cls, tokens), dim=1)
-    pos = torch.cat((-pos.new_ones(len(cls), 1, 2), pos), dim=1)
-    return tokens, pos
-
-
-class Pow3RWrapper(torch.nn.Module):
-    def __init__(
-        self,
-        name,
-        ckpt_path,
-        geometric_input_config,
-        **kwargs,
-    ):
-        super().__init__()
-        self.name = name
-        self.ckpt_path = ckpt_path
-        self.geometric_input_config = geometric_input_config
-
-        # Init the model and load the checkpoint
-        print(f"Loading checkpoint from {self.ckpt_path} ...")
-        ckpt = torch.load(self.ckpt_path, map_location="cpu", weights_only=False)
-        model = ckpt["definition"]
-        print(f"Creating model = {model}")
-        self.model = eval(model)
-        print(self.model.load_state_dict(ckpt["weights"]))
-
-    def forward(self, views):
-        """
-        Forward pass wrapper for Pow3R.
-
-        Assumption:
-        - The number of input views is 2.
-
-        Args:
-            views (List[dict]): List of dictionaries containing the input views' images and instance information.
-                                Length of the list should be 2.
-                                Each dictionary should contain the following keys:
-                                    "img" (tensor): Image tensor of shape (B, C, H, W).
-                                    "data_norm_type" (list): ["dust3r"]
-                                Optionally, each dictionary can also contain the following keys for the respective optional geometric inputs:
-                                    "camera_intrinsics" (tensor): Camera intrinsics. Tensor of shape (B, 3, 3).
-                                    "camera_pose" (tensor): Camera pose. Tensor of shape (B, 4, 4). Camera pose is opencv (RDF) cam2world transformation.
-                                    "depthmap" (tensor): Z Depth map. Tensor of shape (B, H, W, 1).
-
-        Returns:
-            List[dict]: A list containing the final outputs for the two views. Length of the list will be 2.
-        """
-        # Check that the number of input views is 2
-        assert len(views) == 2, "Pow3R requires 2 input views."
-
-        # Check the data norm type
-        data_norm_type = views[0]["data_norm_type"][0]
-        assert data_norm_type == "dust3r", (
-            "Pow3R expects a normalized image with the DUSt3R normalization scheme applied"
-        )
-
-        # Get the batch size per view, device and two views
-        batch_size_per_view = views[0]["img"].shape[0]
-        device = views[0]["img"].device
-        view1, view2 = views
-
-        # Decide if we need to use the geometric inputs
-        if torch.rand(1, device=device) < self.geometric_input_config["overall_prob"]:
-            # Decide if we need to use the camera intrinsics
-            if (
-                torch.rand(1, device=device)
-                < self.geometric_input_config["ray_dirs_prob"]
-            ):
-                add_intrinsics(view1, view1.get("camera_intrinsics"))
-                add_intrinsics(view2, view2.get("camera_intrinsics"))
-
-            # Decide if we need to use the depth map
-            if torch.rand(1, device=device) < self.geometric_input_config["depth_prob"]:
-                depthmap1 = view1.get("depthmap")
-                depthmap2 = view2.get("depthmap")
-                if depthmap1 is not None:
-                    depthmap1 = depthmap1.squeeze(-1).to(device)
-                if depthmap2 is not None:
-                    depthmap2 = depthmap2.squeeze(-1).to(device)
-                add_depth(view1, depthmap1)
-                add_depth(view2, depthmap2)
-
-            # Decide if we need to use the camera pose
-            if torch.rand(1, device=device) < self.geometric_input_config["cam_prob"]:
-                cam1 = view1.get("camera_pose")
-                cam2 = view2.get("camera_pose")
-                add_relpose(view1, cam2_to_world=cam2, cam1_to_world=cam1)
-                add_relpose(view2, cam2_to_world=cam2, cam1_to_world=cam1)
-
-        # Get the model predictions
-        preds = self.model(view1, view2)
-
-        # Convert the output to MapAnything format
-        with torch.autocast("cuda", enabled=False):
-            res = []
-            for view_idx in range(2):
-                # Get the model predictions for the current view
-                curr_view_pred = preds[view_idx]
-
-                # For the first view
-                if view_idx == 0:
-                    # Get the global frame and camera frame pointmaps
-                    global_pts = curr_view_pred["pts3d"]
-                    cam_pts = curr_view_pred["pts3d"]
-                    conf = curr_view_pred["conf"]
-
-                    # Get the ray directions and depth along ray
-                    depth_along_ray = torch.norm(cam_pts, dim=-1, keepdim=True)
-                    ray_directions = cam_pts / depth_along_ray
-
-                    # Initalize identity camera pose
-                    cam_rot = torch.eye(3, device=device)
-                    cam_quat = mat_to_quat(cam_rot)
-                    cam_trans = torch.zeros(3, device=device)
-                    cam_quat = cam_quat.unsqueeze(0).repeat(batch_size_per_view, 1)
-                    cam_trans = cam_trans.unsqueeze(0).repeat(batch_size_per_view, 1)
-                # For the second view
-                elif view_idx == 1:
-                    # Get the global frame and camera frame pointmaps
-                    pred_global_pts = curr_view_pred["pts3d_in_other_view"]
-                    cam_pts = curr_view_pred["pts3d2"]
-                    conf = (curr_view_pred["conf"] * curr_view_pred["conf2"]).sqrt()
-
-                    # Get the ray directions and depth along ray
-                    depth_along_ray = torch.norm(cam_pts, dim=-1, keepdim=True)
-                    ray_directions = cam_pts / depth_along_ray
-
-                    # Compute the camera pose using the pointmaps
-                    cam_rot, cam_trans, scale = roma.rigid_points_registration(
-                        cam_pts.reshape(batch_size_per_view, -1, 3),
-                        pred_global_pts.reshape(batch_size_per_view, -1, 3),
-                        weights=conf.reshape(batch_size_per_view, -1),
-                        compute_scaling=True,
-                    )
-                    cam_quat = mat_to_quat(cam_rot)
-
-                    # Scale the predicted camera frame pointmap and compute the new global frame pointmap
-                    cam_pts = scale.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) * cam_pts
-                    global_pts = cam_pts.reshape(
-                        batch_size_per_view, -1, 3
-                    ) @ cam_rot.permute(0, 2, 1) + cam_trans.unsqueeze(1)
-                    global_pts = global_pts.view(pred_global_pts.shape)
-
-                # Append the result in MapAnything format
-                res.append(
-                    {
-                        "pts3d": global_pts,
-                        "pts3d_cam": cam_pts,
-                        "ray_directions": ray_directions,
-                        "depth_along_ray": depth_along_ray,
-                        "cam_trans": cam_trans,
-                        "cam_quats": cam_quat,
-                        "conf": conf,
-                    }
-                )
-
-        return res
-
-
-class Pow3RBAWrapper(torch.nn.Module):
-    def __init__(
-        self,
-        name,
-        ckpt_path,
-        geometric_input_config,
-        scene_graph="complete",
-        inference_batch_size=32,
-        global_optim_schedule="cosine",
-        global_optim_lr=0.01,
-        global_optim_niter=300,
-        **kwargs,
-    ):
-        super().__init__()
-        self.name = name
-        self.ckpt_path = ckpt_path
-        self.geometric_input_config = geometric_input_config
-        self.scene_graph = scene_graph
-        self.inference_batch_size = inference_batch_size
-        self.global_optim_schedule = global_optim_schedule
-        self.global_optim_lr = global_optim_lr
-        self.global_optim_niter = global_optim_niter
-
-        # Init the model and load the checkpoint
-        print(f"Loading checkpoint from {self.ckpt_path} ...")
-        ckpt = torch.load(self.ckpt_path, map_location="cpu", weights_only=False)
-        model = ckpt["definition"]
-        print(f"Creating model = {model}")
-        self.model = eval(model)
-        print(self.model.load_state_dict(ckpt["weights"]))
-
-        # Init the global aligner mode
-        self.global_aligner_mode = GlobalAlignerMode.PointCloudOptimizer
-
-    def infer_two_views(self, views):
-        """
-        Wrapper for Pow3R 2-View inference.
-
-        Assumption:
-        - The number of input views is 2.
-
-        Args:
-            views (List[dict]): List of dictionaries containing the input views' images and instance information.
-                                Length of the list should be 2.
-                                Each dictionary should contain the following keys:
-                                    "img" (tensor): Image tensor of shape (B, C, H, W).
-                                    "data_norm_type" (list): ["dust3r"]
-                                Optionally, each dictionary can also contain the following keys for the respective optional geometric inputs:
-                                    "camera_intrinsics" (tensor): Camera intrinsics. Tensor of shape (B, 3, 3).
-                                    "camera_pose" (tensor): Camera pose. Tensor of shape (B, 4, 4). Camera pose is opencv (RDF) cam2world transformation.
-                                    "depthmap" (tensor): Z Depth map. Tensor of shape (B, H, W, 1).
-
-        Returns:
-            List[dict]: A list containing the final outputs for the two views. Length of the list will be 2.
-        """
-        # Check that the number of input views is 2
-        assert len(views) == 2, "Pow3R requires 2 input views."
-
-        # Check the data norm type
-        data_norm_type = views[0]["data_norm_type"][0]
-        assert data_norm_type == "dust3r", (
-            "Pow3R expects a normalized image with the DUSt3R normalization scheme applied"
-        )
-
-        # Get the device and two views
-        device = views[0]["img"].device
-        view1, view2 = views
-
-        # Decide if we need to use the geometric inputs
-        if torch.rand(1, device=device) < self.geometric_input_config["overall_prob"]:
-            # Decide if we need to use the camera intrinsics
-            if (
-                torch.rand(1, device=device)
-                < self.geometric_input_config["ray_dirs_prob"]
-            ):
-                add_intrinsics(view1, view1.get("camera_intrinsics"))
-                add_intrinsics(view2, view2.get("camera_intrinsics"))
-
-            # Decide if we need to use the depth map
-            if torch.rand(1, device=device) < self.geometric_input_config["depth_prob"]:
-                depthmap1 = view1.get("depthmap")
-                depthmap2 = view2.get("depthmap")
-                if depthmap1 is not None:
-                    depthmap1 = depthmap1.squeeze(-1).to(device)
-                if depthmap2 is not None:
-                    depthmap2 = depthmap2.squeeze(-1).to(device)
-                add_depth(view1, depthmap1)
-                add_depth(view2, depthmap2)
-
-            # Decide if we need to use the camera pose
-            if torch.rand(1, device=device) < self.geometric_input_config["cam_prob"]:
-                cam1 = view1.get("camera_pose")
-                cam2 = view2.get("camera_pose")
-                add_relpose(view1, cam2_to_world=cam2, cam1_to_world=cam1)
-                add_relpose(view2, cam2_to_world=cam2, cam1_to_world=cam1)
-
-        # Get the model predictions
-        preds = self.model(view1, view2)
-
-        return preds
-
-    def loss_of_one_batch(self, batch, device):
-        """
-        Compute prediction for two views.
-        """
-        view1, view2 = batch
-        ignore_keys = set(
-            [
-                "dataset",
-                "label",
-                "instance",
-                "idx",
-                "true_shape",
-                "rng",
-                "name",
-                "data_norm_type",
-            ]
-        )
-        for view in batch:
-            for name in view.keys():  # pseudo_focal
-                if name in ignore_keys:
-                    continue
-                view[name] = view[name].to(device, non_blocking=True)
-
-        pred1, pred2 = self.infer_two_views([view1, view2])
-
-        result = dict(view1=view1, view2=view2, pred1=pred1, pred2=pred2)
-
-        return result
-
-    @torch.no_grad()
-    def inference(self, pairs, device, verbose=False):
-        """
-        Wrapper for multi-pair inference using Pow3R.
-        """
-        if verbose:
-            print(f">> Inference with model on {len(pairs)} image pairs")
-        result = []
-
-        multiple_shapes = not (check_if_same_size(pairs))
-        if multiple_shapes:
-            self.inference_batch_size = 1
-
-        for i in tqdm.trange(
-            0, len(pairs), self.inference_batch_size, disable=not verbose
-        ):
-            res = self.loss_of_one_batch(
-                collate_with_cat(pairs[i : i + self.inference_batch_size]), device
-            )
-            result.append(to_cpu(res))
-
-        result = collate_with_cat(result, lists=multiple_shapes)
-
-        return result
-
-    def forward(self, views):
-        """
-        Forward pass wrapper for Pow3R using the global aligner.
-
-        Assumption:
-        - The batch size of input views is 1.
-
-        Args:
-            views (List[dict]): List of dictionaries containing the input views' images and instance information.
-                                Each dictionary should contain the following keys, where B is the batch size and is 1:
-                                    "img" (tensor): Image tensor of shape (B, C, H, W).
-                                    "data_norm_type" (list): ["dust3r"]
-
-        Returns:
-            List[dict]: A list containing the final outputs for the input views.
-        """
-        # Check the batch size of input views
-        batch_size_per_view, _, height, width = views[0]["img"].shape
-        device = views[0]["img"].device
-        num_views = len(views)
-        assert batch_size_per_view == 1, (
-            f"Batch size of input views should be 1, but got {batch_size_per_view}."
-        )
-
-        # Check the data norm type
-        data_norm_type = views[0]["data_norm_type"][0]
-        assert data_norm_type == "dust3r", (
-            "Pow3R-BA expects a normalized image with the DUSt3R normalization scheme applied"
-        )
-
-        # Convert the input views to the expected input format
-        images = []
-        for view in views:
-            images.append(
-                dict(
-                    img=view["img"],
-                    camera_intrinsics=view["camera_intrinsics"],
-                    depthmap=view["depthmap"],
-                    camera_pose=view["camera_pose"],
-                    data_norm_type=view["data_norm_type"],
-                    true_shape=view["true_shape"],
-                    idx=len(images),
-                    instance=str(len(images)),
-                )
-            )
-
-        # Make image pairs and run inference pair-wise
-        pairs = make_pairs(
-            images, scene_graph=self.scene_graph, prefilter=None, symmetrize=True
-        )
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore", category=FutureWarning)
-            output = self.inference(
-                pairs,
-                device,
-                verbose=False,
-            )
-
-        # Global optimization
-        with torch.enable_grad():
-            scene = global_aligner(
-                output, device=device, mode=self.global_aligner_mode, verbose=False
-            )
-            _ = scene.compute_global_alignment(
-                init="mst",
-                niter=self.global_optim_niter,
-                schedule=self.global_optim_schedule,
-                lr=self.global_optim_lr,
-            )
-
-        # Make sure scene is not None
-        if scene is None:
-            raise RuntimeError("Global optimization failed.")
-
-        # Get the predictions
-        intrinsics = scene.get_intrinsics()
-        c2w_poses = scene.get_im_poses()
-        depths = scene.get_depthmaps()
-
-        # Convert the output to the MapAnything format
-        with torch.autocast("cuda", enabled=False):
-            res = []
-            for view_idx in range(num_views):
-                # Get the current view predictions
-                curr_view_intrinsic = intrinsics[view_idx].unsqueeze(0)
-                curr_view_pose = c2w_poses[view_idx].unsqueeze(0)
-                curr_view_depth_z = depths[view_idx].unsqueeze(0)
-
-                # Convert the pose to quaternions and translation
-                curr_view_cam_translations = curr_view_pose[..., :3, 3]
-                curr_view_cam_quats = mat_to_quat(curr_view_pose[..., :3, :3])
-
-                # Get the camera frame pointmaps
-                curr_view_pts3d_cam, _ = depthmap_to_camera_frame(
-                    curr_view_depth_z, curr_view_intrinsic
-                )
-
-                # Convert the z depth to depth along ray
-                curr_view_depth_along_ray = convert_z_depth_to_depth_along_ray(
-                    curr_view_depth_z, curr_view_intrinsic
-                )
-                curr_view_depth_along_ray = curr_view_depth_along_ray.unsqueeze(-1)
-
-                # Get the ray directions on the unit sphere in the camera frame
-                _, curr_view_ray_dirs = get_rays_in_camera_frame(
-                    curr_view_intrinsic, height, width, normalize_to_unit_sphere=True
-                )
-
-                # Get the pointmaps
-                curr_view_pts3d = (
-                    convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap(
-                        curr_view_ray_dirs,
-                        curr_view_depth_along_ray,
-                        curr_view_cam_translations,
-                        curr_view_cam_quats,
-                    )
-                )
-
-                # Append the outputs to the result list
-                res.append(
-                    {
-                        "pts3d": curr_view_pts3d,
-                        "pts3d_cam": curr_view_pts3d_cam,
-                        "ray_directions": curr_view_ray_dirs,
-                        "depth_along_ray": curr_view_depth_along_ray,
-                        "cam_trans": curr_view_cam_translations,
-                        "cam_quats": curr_view_cam_quats,
-                    }
-                )
-
-        return res

mapanything/models/external/vggt/__init__.py

@@ -1,186 +0,0 @@
-"""
-Inference wrapper for VGGT
-"""
-
-import torch
-
-from mapanything.models.external.vggt.models.vggt import VGGT
-from mapanything.models.external.vggt.utils.geometry import closed_form_inverse_se3
-from mapanything.models.external.vggt.utils.pose_enc import pose_encoding_to_extri_intri
-from mapanything.models.external.vggt.utils.rotation import mat_to_quat
-from mapanything.utils.geometry import (
-    convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap,
-    convert_z_depth_to_depth_along_ray,
-    depthmap_to_camera_frame,
-    get_rays_in_camera_frame,
-)
-
-
-class VGGTWrapper(torch.nn.Module):
-    def __init__(
-        self,
-        name,
-        torch_hub_force_reload,
-        load_pretrained_weights=True,
-        depth=24,
-        num_heads=16,
-        intermediate_layer_idx=[4, 11, 17, 23],
-        load_custom_ckpt=False,
-        custom_ckpt_path=None,
-    ):
-        super().__init__()
-        self.name = name
-        self.torch_hub_force_reload = torch_hub_force_reload
-        self.load_custom_ckpt = load_custom_ckpt
-        self.custom_ckpt_path = custom_ckpt_path
-
-        if load_pretrained_weights:
-            # Load pre-trained weights
-            if not torch_hub_force_reload:
-                # Initialize the 1B VGGT model from huggingface hub cache
-                print("Loading facebook/VGGT-1B from huggingface cache ...")
-                self.model = VGGT.from_pretrained(
-                    "facebook/VGGT-1B",
-                )
-            else:
-                # Initialize the 1B VGGT model
-                print("Re-downloading facebook/VGGT-1B ...")
-                self.model = VGGT.from_pretrained(
-                    "facebook/VGGT-1B", force_download=True
-                )
-        else:
-            # Load the VGGT class
-            self.model = VGGT(
-                depth=depth,
-                num_heads=num_heads,
-                intermediate_layer_idx=intermediate_layer_idx,
-            )
-
-        # Get the dtype for VGGT inference
-        # bfloat16 is supported on Ampere GPUs (Compute Capability 8.0+)
-        self.dtype = (
-            torch.bfloat16
-            if torch.cuda.get_device_capability()[0] >= 8
-            else torch.float16
-        )
-
-        # Load custom checkpoint if requested
-        if self.load_custom_ckpt:
-            print(f"Loading checkpoint from {self.custom_ckpt_path} ...")
-            assert self.custom_ckpt_path is not None, (
-                "custom_ckpt_path must be provided if load_custom_ckpt is set to True"
-            )
-            custom_ckpt = torch.load(self.custom_ckpt_path, weights_only=False)
-            print(self.model.load_state_dict(custom_ckpt, strict=True))
-            del custom_ckpt  # in case it occupies memory
-
-    def forward(self, views):
-        """
-        Forward pass wrapper for VGGT
-
-        Assumption:
-        - All the input views have the same image shape.
-
-        Args:
-            views (List[dict]): List of dictionaries containing the input views' images and instance information.
-                                Each dictionary should contain the following keys:
-                                    "img" (tensor): Image tensor of shape (B, C, H, W).
-                                    "data_norm_type" (list): ["identity"]
-
-        Returns:
-            List[dict]: A list containing the final outputs for all N views.
-        """
-        # Get input shape of the images, number of views, and batch size per view
-        batch_size_per_view, _, height, width = views[0]["img"].shape
-        num_views = len(views)
-
-        # Check the data norm type
-        # VGGT expects a normalized image but without the DINOv2 mean and std applied ("identity")
-        data_norm_type = views[0]["data_norm_type"][0]
-        assert data_norm_type == "identity", (
-            "VGGT expects a normalized image but without the DINOv2 mean and std applied"
-        )
-
-        # Concatenate the images to create a single (B, V, C, H, W) tensor
-        img_list = [view["img"] for view in views]
-        images = torch.stack(img_list, dim=1)
-
-        # Run the VGGT aggregator
-        with torch.autocast("cuda", dtype=self.dtype):
-            aggregated_tokens_list, ps_idx = self.model.aggregator(images)
-
-        # Run the Camera + Pose Branch of VGGT
-        with torch.autocast("cuda", enabled=False):
-            # Predict Cameras
-            pose_enc = self.model.camera_head(aggregated_tokens_list)[-1]
-            # Extrinsic and intrinsic matrices, following OpenCV convention (camera from world)
-            # Extrinsics Shape: (B, V, 3, 4)
-            # Intrinsics Shape: (B, V, 3, 3)
-            extrinsic, intrinsic = pose_encoding_to_extri_intri(
-                pose_enc, images.shape[-2:]
-            )
-
-            # Predict Depth Maps
-            # Depth Shape: (B, V, H, W, 1)
-            # Depth Confidence Shape: (B, V, H, W)
-            depth_map, depth_conf = self.model.depth_head(
-                aggregated_tokens_list, images, ps_idx
-            )
-
-            # Convert the output to MapAnything format
-            res = []
-            for view_idx in range(num_views):
-                # Get the extrinsics, intrinsics, depth map for the current view
-                curr_view_extrinsic = extrinsic[:, view_idx, ...]
-                curr_view_extrinsic = closed_form_inverse_se3(
-                    curr_view_extrinsic
-                )  # Convert to cam2world
-                curr_view_intrinsic = intrinsic[:, view_idx, ...]
-                curr_view_depth_z = depth_map[:, view_idx, ...]
-                curr_view_depth_z = curr_view_depth_z.squeeze(-1)
-                curr_view_confidence = depth_conf[:, view_idx, ...]
-
-                # Get the camera frame pointmaps
-                curr_view_pts3d_cam, _ = depthmap_to_camera_frame(
-                    curr_view_depth_z, curr_view_intrinsic
-                )
-
-                # Convert the extrinsics to quaternions and translations
-                curr_view_cam_translations = curr_view_extrinsic[..., :3, 3]
-                curr_view_cam_quats = mat_to_quat(curr_view_extrinsic[..., :3, :3])
-
-                # Convert the z depth to depth along ray
-                curr_view_depth_along_ray = convert_z_depth_to_depth_along_ray(
-                    curr_view_depth_z, curr_view_intrinsic
-                )
-                curr_view_depth_along_ray = curr_view_depth_along_ray.unsqueeze(-1)
-
-                # Get the ray directions on the unit sphere in the camera frame
-                _, curr_view_ray_dirs = get_rays_in_camera_frame(
-                    curr_view_intrinsic, height, width, normalize_to_unit_sphere=True
-                )
-
-                # Get the pointmaps
-                curr_view_pts3d = (
-                    convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap(
-                        curr_view_ray_dirs,
-                        curr_view_depth_along_ray,
-                        curr_view_cam_translations,
-                        curr_view_cam_quats,
-                    )
-                )
-
-                # Append the outputs to the result list
-                res.append(
-                    {
-                        "pts3d": curr_view_pts3d,
-                        "pts3d_cam": curr_view_pts3d_cam,
-                        "ray_directions": curr_view_ray_dirs,
-                        "depth_along_ray": curr_view_depth_along_ray,
-                        "cam_trans": curr_view_cam_translations,
-                        "cam_quats": curr_view_cam_quats,
-                        "conf": curr_view_confidence,
-                    }
-                )
-
-        return res

mapanything/models/external/vggt/heads/camera_head.py

@@ -1,167 +0,0 @@
-# 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 mapanything.models.external.vggt.heads.head_act import activate_pose
-from mapanything.models.external.vggt.layers import Mlp
-from mapanything.models.external.vggt.layers.block import Block
-
-
-class CameraHead(nn.Module):
-    """
-    CameraHead predicts camera parameters from token representations using iterative refinement.
-
-    It applies a series of transformer blocks (the "trunk") to dedicated camera tokens.
-    """
-
-    def __init__(
-        self,
-        dim_in: int = 2048,
-        trunk_depth: int = 4,
-        pose_encoding_type: str = "absT_quaR_FoV",
-        num_heads: int = 16,
-        mlp_ratio: int = 4,
-        init_values: float = 0.01,
-        trans_act: str = "linear",
-        quat_act: str = "linear",
-        fl_act: str = "relu",  # Field of view activations: ensures FOV values are positive.
-    ):
-        super().__init__()
-
-        if pose_encoding_type == "absT_quaR_FoV":
-            self.target_dim = 9
-        else:
-            raise ValueError(f"Unsupported camera encoding type: {pose_encoding_type}")
-
-        self.trans_act = trans_act
-        self.quat_act = quat_act
-        self.fl_act = fl_act
-        self.trunk_depth = trunk_depth
-
-        # Build the trunk using a sequence of transformer blocks.
-        self.trunk = nn.Sequential(
-            *[
-                Block(
-                    dim=dim_in,
-                    num_heads=num_heads,
-                    mlp_ratio=mlp_ratio,
-                    init_values=init_values,
-                )
-                for _ in range(trunk_depth)
-            ]
-        )
-
-        # Normalizations for camera token and trunk output.
-        self.token_norm = nn.LayerNorm(dim_in)
-        self.trunk_norm = nn.LayerNorm(dim_in)
-
-        # Learnable empty camera pose token.
-        self.empty_pose_tokens = nn.Parameter(torch.zeros(1, 1, self.target_dim))
-        self.embed_pose = nn.Linear(self.target_dim, dim_in)
-
-        # Module for producing modulation parameters: shift, scale, and a gate.
-        self.poseLN_modulation = nn.Sequential(
-            nn.SiLU(), nn.Linear(dim_in, 3 * dim_in, bias=True)
-        )
-
-        # Adaptive layer normalization without affine parameters.
-        self.adaln_norm = nn.LayerNorm(dim_in, elementwise_affine=False, eps=1e-6)
-        self.pose_branch = Mlp(
-            in_features=dim_in,
-            hidden_features=dim_in // 2,
-            out_features=self.target_dim,
-            drop=0,
-        )
-
-    def forward(self, aggregated_tokens_list: list, num_iterations: int = 4) -> list:
-        """
-        Forward pass to predict camera parameters.
-
-        Args:
-            aggregated_tokens_list (list): List of token tensors from the network;
-                the last tensor is used for prediction.
-            num_iterations (int, optional): Number of iterative refinement steps. Defaults to 4.
-
-        Returns:
-            list: A list of predicted camera encodings (post-activation) from each iteration.
-        """
-        # Use tokens from the last block for camera prediction.
-        tokens = aggregated_tokens_list[-1]
-
-        # Extract the camera tokens
-        pose_tokens = tokens[:, :, 0]
-        pose_tokens = self.token_norm(pose_tokens)
-
-        pred_pose_enc_list = self.trunk_fn(pose_tokens, num_iterations)
-        return pred_pose_enc_list
-
-    def trunk_fn(self, pose_tokens: torch.Tensor, num_iterations: int) -> list:
-        """
-        Iteratively refine camera pose predictions.
-
-        Args:
-            pose_tokens (torch.Tensor): Normalized camera tokens with shape [B, 1, C].
-            num_iterations (int): Number of refinement iterations.
-
-        Returns:
-            list: List of activated camera encodings from each iteration.
-        """
-        B, S, C = pose_tokens.shape  # S is expected to be 1.
-        pred_pose_enc = None
-        pred_pose_enc_list = []
-
-        for _ in range(num_iterations):
-            # Use a learned empty pose for the first iteration.
-            if pred_pose_enc is None:
-                module_input = self.embed_pose(self.empty_pose_tokens.expand(B, S, -1))
-            else:
-                # Detach the previous prediction to avoid backprop through time.
-                pred_pose_enc = pred_pose_enc.detach()
-                module_input = self.embed_pose(pred_pose_enc)
-
-            # Generate modulation parameters and split them into shift, scale, and gate components.
-            shift_msa, scale_msa, gate_msa = self.poseLN_modulation(module_input).chunk(
-                3, dim=-1
-            )
-
-            # Adaptive layer normalization and modulation.
-            pose_tokens_modulated = gate_msa * modulate(
-                self.adaln_norm(pose_tokens), shift_msa, scale_msa
-            )
-            pose_tokens_modulated = pose_tokens_modulated + pose_tokens
-
-            pose_tokens_modulated = self.trunk(pose_tokens_modulated)
-            # Compute the delta update for the pose encoding.
-            pred_pose_enc_delta = self.pose_branch(
-                self.trunk_norm(pose_tokens_modulated)
-            )
-
-            if pred_pose_enc is None:
-                pred_pose_enc = pred_pose_enc_delta
-            else:
-                pred_pose_enc = pred_pose_enc + pred_pose_enc_delta
-
-            # Apply final activation functions for translation, quaternion, and field-of-view.
-            activated_pose = activate_pose(
-                pred_pose_enc,
-                trans_act=self.trans_act,
-                quat_act=self.quat_act,
-                fl_act=self.fl_act,
-            )
-            pred_pose_enc_list.append(activated_pose)
-
-        return pred_pose_enc_list
-
-
-def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
-    """
-    Modulate the input tensor using scaling and shifting parameters.
-    """
-    # modified from https://github.com/facebookresearch/DiT/blob/796c29e532f47bba17c5b9c5eb39b9354b8b7c64/models.py#L19
-    return x * (1 + scale) + shift

mapanything/models/external/vggt/heads/dpt_head.py

@@ -1,600 +0,0 @@
-# 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.
-
-
-# Inspired by https://github.com/DepthAnything/Depth-Anything-V2
-
-
-from typing import List, Tuple, Union
-
-import torch
-import torch.nn as nn
-
-from .head_act import activate_head
-from .utils import create_uv_grid, position_grid_to_embed
-
-
-class DPTHead(nn.Module):
-    """
-    DPT  Head for dense prediction tasks.
-
-    This implementation follows the architecture described in "Vision Transformers for Dense Prediction"
-    (https://arxiv.org/abs/2103.13413). The DPT head processes features from a vision transformer
-    backbone and produces dense predictions by fusing multi-scale features.
-
-    Args:
-        dim_in (int): Input dimension (channels).
-        patch_size (int, optional): Patch size. Default is 14.
-        output_dim (int, optional): Number of output channels. Default is 4.
-        activation (str, optional): Activation type. Default is "inv_log".
-        conf_activation (str, optional): Confidence activation type. Default is "expp1".
-        features (int, optional): Feature channels for intermediate representations. Default is 256.
-        out_channels (List[int], optional): Output channels for each intermediate layer.
-        intermediate_layer_idx (List[int], optional): Indices of layers from aggregated tokens used for DPT.
-        pos_embed (bool, optional): Whether to use positional embedding. Default is True.
-        feature_only (bool, optional): If True, return features only without the last several layers and activation head. Default is False.
-        down_ratio (int, optional): Downscaling factor for the output resolution. Default is 1.
-    """
-
-    def __init__(
-        self,
-        dim_in: int,
-        patch_size: int = 14,
-        output_dim: int = 4,
-        activation: str = "inv_log",
-        conf_activation: str = "expp1",
-        features: int = 256,
-        out_channels: List[int] = [256, 512, 1024, 1024],
-        intermediate_layer_idx: List[int] = [4, 11, 17, 23],
-        pos_embed: bool = True,
-        feature_only: bool = False,
-        down_ratio: int = 1,
-    ) -> None:
-        super(DPTHead, self).__init__()
-        self.patch_size = patch_size
-        self.activation = activation
-        self.conf_activation = conf_activation
-        self.pos_embed = pos_embed
-        self.feature_only = feature_only
-        self.down_ratio = down_ratio
-        self.intermediate_layer_idx = intermediate_layer_idx
-
-        self.norm = nn.LayerNorm(dim_in)
-
-        # Projection layers for each output channel from tokens.
-        self.projects = nn.ModuleList(
-            [
-                nn.Conv2d(
-                    in_channels=dim_in,
-                    out_channels=oc,
-                    kernel_size=1,
-                    stride=1,
-                    padding=0,
-                )
-                for oc in out_channels
-            ]
-        )
-
-        # Resize layers for upsampling feature maps.
-        self.resize_layers = nn.ModuleList(
-            [
-                nn.ConvTranspose2d(
-                    in_channels=out_channels[0],
-                    out_channels=out_channels[0],
-                    kernel_size=4,
-                    stride=4,
-                    padding=0,
-                ),
-                nn.ConvTranspose2d(
-                    in_channels=out_channels[1],
-                    out_channels=out_channels[1],
-                    kernel_size=2,
-                    stride=2,
-                    padding=0,
-                ),
-                nn.Identity(),
-                nn.Conv2d(
-                    in_channels=out_channels[3],
-                    out_channels=out_channels[3],
-                    kernel_size=3,
-                    stride=2,
-                    padding=1,
-                ),
-            ]
-        )
-
-        self.scratch = _make_scratch(
-            out_channels,
-            features,
-            expand=False,
-        )
-
-        # Attach additional modules to scratch.
-        self.scratch.stem_transpose = None
-        self.scratch.refinenet1 = _make_fusion_block(features)
-        self.scratch.refinenet2 = _make_fusion_block(features)
-        self.scratch.refinenet3 = _make_fusion_block(features)
-        self.scratch.refinenet4 = _make_fusion_block(features, has_residual=False)
-
-        head_features_1 = features
-        head_features_2 = 32
-
-        if feature_only:
-            self.scratch.output_conv1 = nn.Conv2d(
-                head_features_1, head_features_1, kernel_size=3, stride=1, padding=1
-            )
-        else:
-            self.scratch.output_conv1 = nn.Conv2d(
-                head_features_1,
-                head_features_1 // 2,
-                kernel_size=3,
-                stride=1,
-                padding=1,
-            )
-            conv2_in_channels = head_features_1 // 2
-
-            self.scratch.output_conv2 = nn.Sequential(
-                nn.Conv2d(
-                    conv2_in_channels,
-                    head_features_2,
-                    kernel_size=3,
-                    stride=1,
-                    padding=1,
-                ),
-                nn.ReLU(inplace=True),
-                nn.Conv2d(
-                    head_features_2, output_dim, kernel_size=1, stride=1, padding=0
-                ),
-            )
-
-    def forward(
-        self,
-        aggregated_tokens_list: List[torch.Tensor],
-        images: torch.Tensor,
-        patch_start_idx: int,
-        frames_chunk_size: int = 8,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        """
-        Forward pass through the DPT head, supports processing by chunking frames.
-        Args:
-            aggregated_tokens_list (List[Tensor]): List of token tensors from different transformer layers.
-            images (Tensor): Input images with shape [B, S, 3, H, W], in range [0, 1].
-            patch_start_idx (int): Starting index for patch tokens in the token sequence.
-                Used to separate patch tokens from other tokens (e.g., camera or register tokens).
-            frames_chunk_size (int, optional): Number of frames to process in each chunk.
-                If None or larger than S, all frames are processed at once. Default: 8.
-
-        Returns:
-            Tensor or Tuple[Tensor, Tensor]:
-                - If feature_only=True: Feature maps with shape [B, S, C, H, W]
-                - Otherwise: Tuple of (predictions, confidence) both with shape [B, S, 1, H, W]
-        """
-        B, S, _, H, W = images.shape
-
-        # If frames_chunk_size is not specified or greater than S, process all frames at once
-        if frames_chunk_size is None or frames_chunk_size >= S:
-            return self._forward_impl(aggregated_tokens_list, images, patch_start_idx)
-
-        # Otherwise, process frames in chunks to manage memory usage
-        assert frames_chunk_size > 0
-
-        # Process frames in batches
-        all_preds = []
-        all_conf = []
-
-        for frames_start_idx in range(0, S, frames_chunk_size):
-            frames_end_idx = min(frames_start_idx + frames_chunk_size, S)
-
-            # Process batch of frames
-            if self.feature_only:
-                chunk_output = self._forward_impl(
-                    aggregated_tokens_list,
-                    images,
-                    patch_start_idx,
-                    frames_start_idx,
-                    frames_end_idx,
-                )
-                all_preds.append(chunk_output)
-            else:
-                chunk_preds, chunk_conf = self._forward_impl(
-                    aggregated_tokens_list,
-                    images,
-                    patch_start_idx,
-                    frames_start_idx,
-                    frames_end_idx,
-                )
-                all_preds.append(chunk_preds)
-                all_conf.append(chunk_conf)
-
-        # Concatenate results along the sequence dimension
-        if self.feature_only:
-            return torch.cat(all_preds, dim=1)
-        else:
-            return torch.cat(all_preds, dim=1), torch.cat(all_conf, dim=1)
-
-    def _forward_impl(
-        self,
-        aggregated_tokens_list: List[torch.Tensor],
-        images: torch.Tensor,
-        patch_start_idx: int,
-        frames_start_idx: int = None,
-        frames_end_idx: int = None,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        """
-        Implementation of the forward pass through the DPT head.
-
-        This method processes a specific chunk of frames from the sequence.
-
-        Args:
-            aggregated_tokens_list (List[Tensor]): List of token tensors from different transformer layers.
-            images (Tensor): Input images with shape [B, S, 3, H, W].
-            patch_start_idx (int): Starting index for patch tokens.
-            frames_start_idx (int, optional): Starting index for frames to process.
-            frames_end_idx (int, optional): Ending index for frames to process.
-
-        Returns:
-            Tensor or Tuple[Tensor, Tensor]: Feature maps or (predictions, confidence).
-        """
-        if frames_start_idx is not None and frames_end_idx is not None:
-            images = images[:, frames_start_idx:frames_end_idx].contiguous()
-
-        B, S, _, H, W = images.shape
-
-        patch_h, patch_w = H // self.patch_size, W // self.patch_size
-
-        out = []
-        dpt_idx = 0
-
-        for layer_idx in self.intermediate_layer_idx:
-            x = aggregated_tokens_list[layer_idx][:, :, patch_start_idx:]
-
-            # Select frames if processing a chunk
-            if frames_start_idx is not None and frames_end_idx is not None:
-                x = x[:, frames_start_idx:frames_end_idx]
-
-            x = x.reshape(B * S, -1, x.shape[-1])
-
-            x = self.norm(x)
-
-            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
-
-            x = self.projects[dpt_idx](x)
-            if self.pos_embed:
-                x = self._apply_pos_embed(x, W, H)
-            x = self.resize_layers[dpt_idx](x)
-
-            out.append(x)
-            dpt_idx += 1
-
-        # Fuse features from multiple layers.
-        out = self.scratch_forward(out)
-        # Interpolate fused output to match target image resolution.
-        out = custom_interpolate(
-            out,
-            (
-                int(patch_h * self.patch_size / self.down_ratio),
-                int(patch_w * self.patch_size / self.down_ratio),
-            ),
-            mode="bilinear",
-            align_corners=True,
-        )
-
-        if self.pos_embed:
-            out = self._apply_pos_embed(out, W, H)
-
-        if self.feature_only:
-            return out.view(B, S, *out.shape[1:])
-
-        out = self.scratch.output_conv2(out)
-        preds, conf = activate_head(
-            out, activation=self.activation, conf_activation=self.conf_activation
-        )
-
-        preds = preds.view(B, S, *preds.shape[1:])
-        conf = conf.view(B, S, *conf.shape[1:])
-        return preds, conf
-
-    def _apply_pos_embed(
-        self, x: torch.Tensor, W: int, H: int, ratio: float = 0.1
-    ) -> torch.Tensor:
-        """
-        Apply positional embedding to tensor x.
-        """
-        patch_w = x.shape[-1]
-        patch_h = x.shape[-2]
-        pos_embed = create_uv_grid(
-            patch_w, patch_h, aspect_ratio=W / H, dtype=x.dtype, device=x.device
-        )
-        pos_embed = position_grid_to_embed(pos_embed, x.shape[1])
-        pos_embed = pos_embed * ratio
-        pos_embed = pos_embed.permute(2, 0, 1)[None].expand(x.shape[0], -1, -1, -1)
-        return x + pos_embed
-
-    def scratch_forward(self, features: List[torch.Tensor]) -> torch.Tensor:
-        """
-        Forward pass through the fusion blocks.
-
-        Args:
-            features (List[Tensor]): List of feature maps from different layers.
-
-        Returns:
-            Tensor: Fused feature map.
-        """
-        layer_1, layer_2, layer_3, layer_4 = features
-
-        layer_1_rn = self.scratch.layer1_rn(layer_1)
-        layer_2_rn = self.scratch.layer2_rn(layer_2)
-        layer_3_rn = self.scratch.layer3_rn(layer_3)
-        layer_4_rn = self.scratch.layer4_rn(layer_4)
-
-        out = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
-        del layer_4_rn, layer_4
-
-        out = self.scratch.refinenet3(out, layer_3_rn, size=layer_2_rn.shape[2:])
-        del layer_3_rn, layer_3
-
-        out = self.scratch.refinenet2(out, layer_2_rn, size=layer_1_rn.shape[2:])
-        del layer_2_rn, layer_2
-
-        out = self.scratch.refinenet1(out, layer_1_rn)
-        del layer_1_rn, layer_1
-
-        out = self.scratch.output_conv1(out)
-        return out
-
-
-################################################################################
-# Modules
-################################################################################
-
-
-def _make_fusion_block(
-    features: int, size: int = None, has_residual: bool = True, groups: int = 1
-) -> nn.Module:
-    return FeatureFusionBlock(
-        features,
-        nn.ReLU(inplace=True),
-        deconv=False,
-        bn=False,
-        expand=False,
-        align_corners=True,
-        size=size,
-        has_residual=has_residual,
-        groups=groups,
-    )
-
-
-def _make_scratch(
-    in_shape: List[int], out_shape: int, groups: int = 1, expand: bool = False
-) -> nn.Module:
-    scratch = nn.Module()
-    out_shape1 = out_shape
-    out_shape2 = out_shape
-    out_shape3 = out_shape
-    if len(in_shape) >= 4:
-        out_shape4 = out_shape
-
-    if expand:
-        out_shape1 = out_shape
-        out_shape2 = out_shape * 2
-        out_shape3 = out_shape * 4
-        if len(in_shape) >= 4:
-            out_shape4 = out_shape * 8
-
-    scratch.layer1_rn = nn.Conv2d(
-        in_shape[0],
-        out_shape1,
-        kernel_size=3,
-        stride=1,
-        padding=1,
-        bias=False,
-        groups=groups,
-    )
-    scratch.layer2_rn = nn.Conv2d(
-        in_shape[1],
-        out_shape2,
-        kernel_size=3,
-        stride=1,
-        padding=1,
-        bias=False,
-        groups=groups,
-    )
-    scratch.layer3_rn = nn.Conv2d(
-        in_shape[2],
-        out_shape3,
-        kernel_size=3,
-        stride=1,
-        padding=1,
-        bias=False,
-        groups=groups,
-    )
-    if len(in_shape) >= 4:
-        scratch.layer4_rn = nn.Conv2d(
-            in_shape[3],
-            out_shape4,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            bias=False,
-            groups=groups,
-        )
-    return scratch
-
-
-class ResidualConvUnit(nn.Module):
-    """Residual convolution module."""
-
-    def __init__(self, features, activation, bn, groups=1):
-        """Init.
-
-        Args:
-            features (int): number of features
-        """
-        super().__init__()
-
-        self.bn = bn
-        self.groups = groups
-        self.conv1 = nn.Conv2d(
-            features,
-            features,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            bias=True,
-            groups=self.groups,
-        )
-        self.conv2 = nn.Conv2d(
-            features,
-            features,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            bias=True,
-            groups=self.groups,
-        )
-
-        self.norm1 = None
-        self.norm2 = None
-
-        self.activation = activation
-        self.skip_add = nn.quantized.FloatFunctional()
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input
-
-        Returns:
-            tensor: output
-        """
-
-        out = self.activation(x)
-        out = self.conv1(out)
-        if self.norm1 is not None:
-            out = self.norm1(out)
-
-        out = self.activation(out)
-        out = self.conv2(out)
-        if self.norm2 is not None:
-            out = self.norm2(out)
-
-        return self.skip_add.add(out, x)
-
-
-class FeatureFusionBlock(nn.Module):
-    """Feature fusion block."""
-
-    def __init__(
-        self,
-        features,
-        activation,
-        deconv=False,
-        bn=False,
-        expand=False,
-        align_corners=True,
-        size=None,
-        has_residual=True,
-        groups=1,
-    ):
-        """Init.
-
-        Args:
-            features (int): number of features
-        """
-        super(FeatureFusionBlock, self).__init__()
-
-        self.deconv = deconv
-        self.align_corners = align_corners
-        self.groups = groups
-        self.expand = expand
-        out_features = features
-        if self.expand:
-            out_features = features // 2
-
-        self.out_conv = nn.Conv2d(
-            features,
-            out_features,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            bias=True,
-            groups=self.groups,
-        )
-
-        if has_residual:
-            self.resConfUnit1 = ResidualConvUnit(
-                features, activation, bn, groups=self.groups
-            )
-
-        self.has_residual = has_residual
-        self.resConfUnit2 = ResidualConvUnit(
-            features, activation, bn, groups=self.groups
-        )
-
-        self.skip_add = nn.quantized.FloatFunctional()
-        self.size = size
-
-    def forward(self, *xs, size=None):
-        """Forward pass.
-
-        Returns:
-            tensor: output
-        """
-        output = xs[0]
-
-        if self.has_residual:
-            res = self.resConfUnit1(xs[1])
-            output = self.skip_add.add(output, res)
-
-        output = self.resConfUnit2(output)
-
-        if (size is None) and (self.size is None):
-            modifier = {"scale_factor": 2}
-        elif size is None:
-            modifier = {"size": self.size}
-        else:
-            modifier = {"size": size}
-
-        output = custom_interpolate(
-            output, **modifier, mode="bilinear", align_corners=self.align_corners
-        )
-        output = self.out_conv(output)
-
-        return output
-
-
-def custom_interpolate(
-    x: torch.Tensor,
-    size: Tuple[int, int] = None,
-    scale_factor: float = None,
-    mode: str = "bilinear",
-    align_corners: bool = True,
-) -> torch.Tensor:
-    """
-    Custom interpolate to avoid INT_MAX issues in nn.functional.interpolate.
-    """
-    if size is None:
-        size = (int(x.shape[-2] * scale_factor), int(x.shape[-1] * scale_factor))
-
-    INT_MAX = 1610612736
-
-    input_elements = size[0] * size[1] * x.shape[0] * x.shape[1]
-
-    if input_elements > INT_MAX:
-        chunks = torch.chunk(x, chunks=(input_elements // INT_MAX) + 1, dim=0)
-        interpolated_chunks = [
-            nn.functional.interpolate(
-                chunk, size=size, mode=mode, align_corners=align_corners
-            )
-            for chunk in chunks
-        ]
-        x = torch.cat(interpolated_chunks, dim=0)
-        return x.contiguous()
-    else:
-        return nn.functional.interpolate(
-            x, size=size, mode=mode, align_corners=align_corners
-        )

mapanything/models/external/vggt/heads/head_act.py

@@ -1,127 +0,0 @@
-# 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.functional as F
-
-
-def activate_pose(
-    pred_pose_enc, trans_act="linear", quat_act="linear", fl_act="linear"
-):
-    """
-    Activate pose parameters with specified activation functions.
-
-    Args:
-        pred_pose_enc: Tensor containing encoded pose parameters [translation, quaternion, focal length]
-        trans_act: Activation type for translation component
-        quat_act: Activation type for quaternion component
-        fl_act: Activation type for focal length component
-
-    Returns:
-        Activated pose parameters tensor
-    """
-    T = pred_pose_enc[..., :3]
-    quat = pred_pose_enc[..., 3:7]
-    fl = pred_pose_enc[..., 7:]  # or fov
-
-    T = base_pose_act(T, trans_act)
-    quat = base_pose_act(quat, quat_act)
-    fl = base_pose_act(fl, fl_act)  # or fov
-
-    pred_pose_enc = torch.cat([T, quat, fl], dim=-1)
-
-    return pred_pose_enc
-
-
-def base_pose_act(pose_enc, act_type="linear"):
-    """
-    Apply basic activation function to pose parameters.
-
-    Args:
-        pose_enc: Tensor containing encoded pose parameters
-        act_type: Activation type ("linear", "inv_log", "exp", "relu")
-
-    Returns:
-        Activated pose parameters
-    """
-    if act_type == "linear":
-        return pose_enc
-    elif act_type == "inv_log":
-        return inverse_log_transform(pose_enc)
-    elif act_type == "exp":
-        return torch.exp(pose_enc)
-    elif act_type == "relu":
-        return F.relu(pose_enc)
-    else:
-        raise ValueError(f"Unknown act_type: {act_type}")
-
-
-def activate_head(out, activation="norm_exp", conf_activation="expp1"):
-    """
-    Process network output to extract 3D points and confidence values.
-
-    Args:
-        out: Network output tensor (B, C, H, W)
-        activation: Activation type for 3D points
-        conf_activation: Activation type for confidence values
-
-    Returns:
-        Tuple of (3D points tensor, confidence tensor)
-    """
-    # Move channels from last dim to the 4th dimension => (B, H, W, C)
-    fmap = out.permute(0, 2, 3, 1)  # B,H,W,C expected
-
-    # Split into xyz (first C-1 channels) and confidence (last channel)
-    xyz = fmap[:, :, :, :-1]
-    conf = fmap[:, :, :, -1]
-
-    if activation == "norm_exp":
-        d = xyz.norm(dim=-1, keepdim=True).clamp(min=1e-8)
-        xyz_normed = xyz / d
-        pts3d = xyz_normed * torch.expm1(d)
-    elif activation == "norm":
-        pts3d = xyz / xyz.norm(dim=-1, keepdim=True)
-    elif activation == "exp":
-        pts3d = torch.exp(xyz)
-    elif activation == "relu":
-        pts3d = F.relu(xyz)
-    elif activation == "inv_log":
-        pts3d = inverse_log_transform(xyz)
-    elif activation == "xy_inv_log":
-        xy, z = xyz.split([2, 1], dim=-1)
-        z = inverse_log_transform(z)
-        pts3d = torch.cat([xy * z, z], dim=-1)
-    elif activation == "sigmoid":
-        pts3d = torch.sigmoid(xyz)
-    elif activation == "linear":
-        pts3d = xyz
-    else:
-        raise ValueError(f"Unknown activation: {activation}")
-
-    if conf_activation == "expp1":
-        conf_out = 1 + conf.exp()
-    elif conf_activation == "expp0":
-        conf_out = conf.exp()
-    elif conf_activation == "sigmoid":
-        conf_out = torch.sigmoid(conf)
-    else:
-        raise ValueError(f"Unknown conf_activation: {conf_activation}")
-
-    return pts3d, conf_out
-
-
-def inverse_log_transform(y):
-    """
-    Apply inverse log transform: sign(y) * (exp(|y|) - 1)
-
-    Args:
-        y: Input tensor
-
-    Returns:
-        Transformed tensor
-    """
-    return torch.sign(y) * (torch.expm1(torch.abs(y)))

mapanything/models/external/vggt/heads/track_head.py

@@ -1,118 +0,0 @@
-# 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.nn as nn
-
-from .dpt_head import DPTHead
-from .track_modules.base_track_predictor import BaseTrackerPredictor
-
-
-class TrackHead(nn.Module):
-    """
-    Track head that uses DPT head to process tokens and BaseTrackerPredictor for tracking.
-    The tracking is performed iteratively, refining predictions over multiple iterations.
-    """
-
-    def __init__(
-        self,
-        dim_in,
-        patch_size=14,
-        features=128,
-        iters=4,
-        predict_conf=True,
-        stride=2,
-        corr_levels=7,
-        corr_radius=4,
-        hidden_size=384,
-    ):
-        """
-        Initialize the TrackHead module.
-
-        Args:
-            dim_in (int): Input dimension of tokens from the backbone.
-            patch_size (int): Size of image patches used in the vision transformer.
-            features (int): Number of feature channels in the feature extractor output.
-            iters (int): Number of refinement iterations for tracking predictions.
-            predict_conf (bool): Whether to predict confidence scores for tracked points.
-            stride (int): Stride value for the tracker predictor.
-            corr_levels (int): Number of correlation pyramid levels
-            corr_radius (int): Radius for correlation computation, controlling the search area.
-            hidden_size (int): Size of hidden layers in the tracker network.
-        """
-        super().__init__()
-
-        self.patch_size = patch_size
-
-        # Feature extractor based on DPT architecture
-        # Processes tokens into feature maps for tracking
-        self.feature_extractor = DPTHead(
-            dim_in=dim_in,
-            patch_size=patch_size,
-            features=features,
-            feature_only=True,  # Only output features, no activation
-            down_ratio=2,  # Reduces spatial dimensions by factor of 2
-            pos_embed=False,
-        )
-
-        # Tracker module that predicts point trajectories
-        # Takes feature maps and predicts coordinates and visibility
-        self.tracker = BaseTrackerPredictor(
-            latent_dim=features,  # Match the output_dim of feature extractor
-            predict_conf=predict_conf,
-            stride=stride,
-            corr_levels=corr_levels,
-            corr_radius=corr_radius,
-            hidden_size=hidden_size,
-        )
-
-        self.iters = iters
-
-    def forward(
-        self,
-        aggregated_tokens_list,
-        images,
-        patch_start_idx,
-        query_points=None,
-        iters=None,
-    ):
-        """
-        Forward pass of the TrackHead.
-
-        Args:
-            aggregated_tokens_list (list): List of aggregated tokens from the backbone.
-            images (torch.Tensor): Input images of shape (B, S, C, H, W) where:
-                                   B = batch size, S = sequence length.
-            patch_start_idx (int): Starting index for patch tokens.
-            query_points (torch.Tensor, optional): Initial query points to track.
-                                                  If None, points are initialized by the tracker.
-            iters (int, optional): Number of refinement iterations. If None, uses self.iters.
-
-        Returns:
-            tuple:
-                - coord_preds (torch.Tensor): Predicted coordinates for tracked points.
-                - vis_scores (torch.Tensor): Visibility scores for tracked points.
-                - conf_scores (torch.Tensor): Confidence scores for tracked points (if predict_conf=True).
-        """
-        B, S, _, H, W = images.shape
-
-        # Extract features from tokens
-        # feature_maps has shape (B, S, C, H//2, W//2) due to down_ratio=2
-        feature_maps = self.feature_extractor(
-            aggregated_tokens_list, images, patch_start_idx
-        )
-
-        # Use default iterations if not specified
-        if iters is None:
-            iters = self.iters
-
-        # Perform tracking using the extracted features
-        coord_preds, vis_scores, conf_scores = self.tracker(
-            query_points=query_points,
-            fmaps=feature_maps,
-            iters=iters,
-        )
-
-        return coord_preds, vis_scores, conf_scores

mapanything/models/external/vggt/heads/track_modules/__init__.py

@@ -1,5 +0,0 @@
-# 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.

mapanything/models/external/vggt/heads/track_modules/base_track_predictor.py

@@ -1,242 +0,0 @@
-# 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 einops import rearrange
-
-from .blocks import CorrBlock, EfficientUpdateFormer
-from .modules import Mlp
-from .utils import get_2d_embedding, get_2d_sincos_pos_embed, sample_features4d
-
-
-class BaseTrackerPredictor(nn.Module):
-    def __init__(
-        self,
-        stride=1,
-        corr_levels=5,
-        corr_radius=4,
-        latent_dim=128,
-        hidden_size=384,
-        use_spaceatt=True,
-        depth=6,
-        max_scale=518,
-        predict_conf=True,
-    ):
-        super(BaseTrackerPredictor, self).__init__()
-        """
-        The base template to create a track predictor
-
-        Modified from https://github.com/facebookresearch/co-tracker/
-        and https://github.com/facebookresearch/vggsfm
-        """
-
-        self.stride = stride
-        self.latent_dim = latent_dim
-        self.corr_levels = corr_levels
-        self.corr_radius = corr_radius
-        self.hidden_size = hidden_size
-        self.max_scale = max_scale
-        self.predict_conf = predict_conf
-
-        self.flows_emb_dim = latent_dim // 2
-
-        self.corr_mlp = Mlp(
-            in_features=self.corr_levels * (self.corr_radius * 2 + 1) ** 2,
-            hidden_features=self.hidden_size,
-            out_features=self.latent_dim,
-        )
-
-        self.transformer_dim = self.latent_dim + self.latent_dim + self.latent_dim + 4
-
-        self.query_ref_token = nn.Parameter(torch.randn(1, 2, self.transformer_dim))
-
-        space_depth = depth if use_spaceatt else 0
-        time_depth = depth
-
-        self.updateformer = EfficientUpdateFormer(
-            space_depth=space_depth,
-            time_depth=time_depth,
-            input_dim=self.transformer_dim,
-            hidden_size=self.hidden_size,
-            output_dim=self.latent_dim + 2,
-            mlp_ratio=4.0,
-            add_space_attn=use_spaceatt,
-        )
-
-        self.fmap_norm = nn.LayerNorm(self.latent_dim)
-        self.ffeat_norm = nn.GroupNorm(1, self.latent_dim)
-
-        # A linear layer to update track feats at each iteration
-        self.ffeat_updater = nn.Sequential(
-            nn.Linear(self.latent_dim, self.latent_dim), nn.GELU()
-        )
-
-        self.vis_predictor = nn.Sequential(nn.Linear(self.latent_dim, 1))
-
-        if predict_conf:
-            self.conf_predictor = nn.Sequential(nn.Linear(self.latent_dim, 1))
-
-    def forward(
-        self,
-        query_points,
-        fmaps=None,
-        iters=6,
-        return_feat=False,
-        down_ratio=1,
-        apply_sigmoid=True,
-    ):
-        """
-        query_points: B x N x 2, the number of batches, tracks, and xy
-        fmaps: B x S x C x HH x WW, the number of batches, frames, and feature dimension.
-                note HH and WW is the size of feature maps instead of original images
-        """
-        B, N, D = query_points.shape
-        B, S, C, HH, WW = fmaps.shape
-
-        assert D == 2, "Input points must be 2D coordinates"
-
-        # apply a layernorm to fmaps here
-        fmaps = self.fmap_norm(fmaps.permute(0, 1, 3, 4, 2))
-        fmaps = fmaps.permute(0, 1, 4, 2, 3)
-
-        # Scale the input query_points because we may downsample the images
-        # by down_ratio or self.stride
-        # e.g., if a 3x1024x1024 image is processed to a 128x256x256 feature map
-        # its query_points should be query_points/4
-        if down_ratio > 1:
-            query_points = query_points / float(down_ratio)
-
-        query_points = query_points / float(self.stride)
-
-        # Init with coords as the query points
-        # It means the search will start from the position of query points at the reference frames
-        coords = query_points.clone().reshape(B, 1, N, 2).repeat(1, S, 1, 1)
-
-        # Sample/extract the features of the query points in the query frame
-        query_track_feat = sample_features4d(fmaps[:, 0], coords[:, 0])
-
-        # init track feats by query feats
-        track_feats = query_track_feat.unsqueeze(1).repeat(1, S, 1, 1)  # B, S, N, C
-        # back up the init coords
-        coords_backup = coords.clone()
-
-        fcorr_fn = CorrBlock(
-            fmaps, num_levels=self.corr_levels, radius=self.corr_radius
-        )
-
-        coord_preds = []
-
-        # Iterative Refinement
-        for _ in range(iters):
-            # Detach the gradients from the last iteration
-            # (in my experience, not very important for performance)
-            coords = coords.detach()
-
-            fcorrs = fcorr_fn.corr_sample(track_feats, coords)
-
-            corr_dim = fcorrs.shape[3]
-            fcorrs_ = fcorrs.permute(0, 2, 1, 3).reshape(B * N, S, corr_dim)
-            fcorrs_ = self.corr_mlp(fcorrs_)
-
-            # Movement of current coords relative to query points
-            flows = (coords - coords[:, 0:1]).permute(0, 2, 1, 3).reshape(B * N, S, 2)
-
-            flows_emb = get_2d_embedding(flows, self.flows_emb_dim, cat_coords=False)
-
-            # (In my trials, it is also okay to just add the flows_emb instead of concat)
-            flows_emb = torch.cat(
-                [flows_emb, flows / self.max_scale, flows / self.max_scale], dim=-1
-            )
-
-            track_feats_ = track_feats.permute(0, 2, 1, 3).reshape(
-                B * N, S, self.latent_dim
-            )
-
-            # Concatenate them as the input for the transformers
-            transformer_input = torch.cat([flows_emb, fcorrs_, track_feats_], dim=2)
-
-            # 2D positional embed
-            # TODO: this can be much simplified
-            pos_embed = get_2d_sincos_pos_embed(
-                self.transformer_dim, grid_size=(HH, WW)
-            ).to(query_points.device)
-            sampled_pos_emb = sample_features4d(
-                pos_embed.expand(B, -1, -1, -1), coords[:, 0]
-            )
-
-            sampled_pos_emb = rearrange(sampled_pos_emb, "b n c -> (b n) c").unsqueeze(
-                1
-            )
-
-            x = transformer_input + sampled_pos_emb
-
-            # Add the query ref token to the track feats
-            query_ref_token = torch.cat(
-                [
-                    self.query_ref_token[:, 0:1],
-                    self.query_ref_token[:, 1:2].expand(-1, S - 1, -1),
-                ],
-                dim=1,
-            )
-            x = x + query_ref_token.to(x.device).to(x.dtype)
-
-            # B, N, S, C
-            x = rearrange(x, "(b n) s d -> b n s d", b=B)
-
-            # Compute the delta coordinates and delta track features
-            delta, _ = self.updateformer(x)
-
-            # BN, S, C
-            delta = rearrange(delta, " b n s d -> (b n) s d", b=B)
-            delta_coords_ = delta[:, :, :2]
-            delta_feats_ = delta[:, :, 2:]
-
-            track_feats_ = track_feats_.reshape(B * N * S, self.latent_dim)
-            delta_feats_ = delta_feats_.reshape(B * N * S, self.latent_dim)
-
-            # Update the track features
-            track_feats_ = (
-                self.ffeat_updater(self.ffeat_norm(delta_feats_)) + track_feats_
-            )
-
-            track_feats = track_feats_.reshape(B, N, S, self.latent_dim).permute(
-                0, 2, 1, 3
-            )  # BxSxNxC
-
-            # B x S x N x 2
-            coords = coords + delta_coords_.reshape(B, N, S, 2).permute(0, 2, 1, 3)
-
-            # Force coord0 as query
-            # because we assume the query points should not be changed
-            coords[:, 0] = coords_backup[:, 0]
-
-            # The predicted tracks are in the original image scale
-            if down_ratio > 1:
-                coord_preds.append(coords * self.stride * down_ratio)
-            else:
-                coord_preds.append(coords * self.stride)
-
-        # B, S, N
-        vis_e = self.vis_predictor(
-            track_feats.reshape(B * S * N, self.latent_dim)
-        ).reshape(B, S, N)
-        if apply_sigmoid:
-            vis_e = torch.sigmoid(vis_e)
-
-        if self.predict_conf:
-            conf_e = self.conf_predictor(
-                track_feats.reshape(B * S * N, self.latent_dim)
-            ).reshape(B, S, N)
-            if apply_sigmoid:
-                conf_e = torch.sigmoid(conf_e)
-        else:
-            conf_e = None
-
-        if return_feat:
-            return coord_preds, vis_e, track_feats, query_track_feat, conf_e
-        else:
-            return coord_preds, vis_e, conf_e