Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

third_party/vipe/.vscode/settings.json

@@ -0,0 +1,60 @@
+{
+    "[cpp]": {
+        "editor.formatOnSave": true,
+        "editor.defaultFormatter": "ms-vscode.cpptools"
+    },
+    "[cuda-cpp]": {
+        "editor.formatOnSave": true,
+        "editor.defaultFormatter": "ms-vscode.cpptools"
+    },
+    "[python]": {
+        "editor.defaultFormatter": "charliermarsh.ruff",
+        "editor.formatOnSave": true,
+        "editor.codeActionsOnSave": {
+          "source.organizeImports": "explicit"
+        },
+    },
+    "pylint.enabled": false,
+    "files.associations": {
+        "optional": "cpp",
+        "istream": "cpp",
+        "ostream": "cpp",
+        "ratio": "cpp",
+        "system_error": "cpp",
+        "array": "cpp",
+        "functional": "cpp",
+        "tuple": "cpp",
+        "type_traits": "cpp",
+        "utility": "cpp",
+        "variant": "cpp",
+        "charconv": "cpp",
+        "chrono": "cpp",
+        "format": "cpp",
+        "__bit_reference": "cpp",
+        "__node_handle": "cpp",
+        "bitset": "cpp",
+        "deque": "cpp",
+        "limits": "cpp",
+        "locale": "cpp",
+        "vector": "cpp",
+        "complex": "cpp",
+        "string_view": "cpp",
+        "initializer_list": "cpp",
+        "ranges": "cpp",
+        "span": "cpp",
+        "__hash_table": "cpp",
+        "__split_buffer": "cpp",
+        "__tree": "cpp",
+        "list": "cpp",
+        "map": "cpp",
+        "queue": "cpp",
+        "set": "cpp",
+        "stack": "cpp",
+        "string": "cpp",
+        "unordered_map": "cpp",
+        "unordered_set": "cpp",
+        "iostream": "cpp",
+        "dense": "cpp",
+        "*.tcc": "cpp"
+    },
+}

third_party/vipe/CONTRIBUTING.md

@@ -0,0 +1,52 @@
+# ViPE Contribution Rules
+
+## Coding Guidelines
+
+- Please follow the existing conventions in the relevant file, submodule, module, and project when you add new code or when you extend/fix existing functionality.
+
+- Make sure that you can contribute your work to open source (no license and/or patent conflict is introduced by your code). You will need to [`sign`](#signing-your-work) your commit.
+
+- Thanks in advance for your patience as we review your contributions; we do appreciate them!
+
+## Signing Your Work
+
+* We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
+
+  * Any contribution which contains commits that are not Signed-Off will not be accepted.
+
+* To sign off on a commit you simply use the `--signoff` (or `-s`) option when committing your changes:
+  ```bash
+  $ git commit -s -m "Add cool feature."
+  ```
+  This will append the following to your commit message:
+  ```
+  Signed-off-by: Your Name <your@email.com>
+  ```
+
+* Full text of the DCO:
+
+  ```
+    Developer Certificate of Origin
+    Version 1.1
+    
+    Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
+    1 Letterman Drive
+    Suite D4700
+    San Francisco, CA, 94129
+    
+    Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed.
+  ```
+
+  ```
+    Developer's Certificate of Origin 1.1
+    
+    By making a contribution to this project, I certify that:
+    
+    (a) The contribution was created in whole or in part by me and I have the right to submit it under the open source license indicated in the file; or
+    
+    (b) The contribution is based upon previous work that, to the best of my knowledge, is covered under an appropriate open source license and I have the right under that license to submit that work with modifications, whether created in whole or in part by me, under the same open source license (unless I am permitted to submit under a different license), as indicated in the file; or
+    
+    (c) The contribution was provided directly to me by some other person who certified (a), (b) or (c) and I have not modified it.
+    
+    (d) I understand and agree that this project and the contribution are public and that a record of the contribution (including all personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with this project or the open source license(s) involved.
+  ```

third_party/vipe/LICENSE

@@ -0,0 +1,13 @@
+Copyright 2025 NVIDIA CORPORATION & AFFILIATES.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

third_party/vipe/README.md

@@ -0,0 +1,153 @@
+# ViPE: Video Pose Engine for Geometric 3D Perception
+
+<p align="center">
+  <img src="assets/teaser.gif" alt="teaser"/>
+</p>
+
+**TL;DR: ViPE is a useful open-source spatial AI tool for annotating camera poses and dense depth maps from raw videos!**
+
+**Contributors**: NVIDIA (Spatial Intelligence Lab, Dynamic Vision Lab, NVIDIA Issac, NVIDIA Research).
+
+**Full Abstract**: Accurate 3D geometric perception is an important prerequisite for a wide range of spatial AI systems. While state-of-the-art methods depend on large-scale training data, acquiring consistent and precise 3D annotations from in-the-wild videos remains a key challenge. In this work, we introduce ViPE, a handy and versatile video processing engine designed to bridge this gap. ViPE efficiently estimates camera intrinsics, camera motion, and dense, near-metric depth maps from unconstrained raw videos. It is robust to diverse scenarios, including dynamic selfie videos, cinematic shots, or dashcams, and supports various camera models such as pinhole, wide-angle, and 360° panoramas. 
+We use ViPE to annotate a large-scale collection of videos. This collection includes around 100K real-world internet videos, 1M high-quality AI-generated videos, and 2K panoramic videos, totaling approximately 96M frames -- all annotated with accurate camera poses and dense depth maps. We open source ViPE and the annotated dataset with the hope to accelerate the development of spatial AI systems.
+
+**[Technical Whitepaper](https://research.nvidia.com/labs/toronto-ai/vipe/assets/paper.pdf), [Project Page](https://research.nvidia.com/labs/toronto-ai/vipe), [Dataset](#downloading-the-dataset)**
+
+## Installation
+
+To ensure the reproducibility, we recommend creating the runtime environment using [conda](https://docs.conda.io/projects/conda/en/latest/user-guide/getting-started.html).
+
+```bash
+# Create a new conda environment and install 3rd-party dependencies
+conda env create -f envs/base.yml
+conda activate vipe
+# You can switch to your own PyPI index if you want.
+pip install -r envs/requirements.txt --extra-index-url https://download.pytorch.org/whl/cu128
+
+# Build the project and install it into the current environment
+# Omit the -e flag to install the project as a regular package
+pip install --no-build-isolation -e .
+```
+
+## Usage
+
+### Using the ViPE CLI
+
+Once the python package is installed, you can use the `vipe` CLI to process raw videos in mp4 format.
+
+```bash
+# Replace YOUR_VIDEO.mp4 with the path to your video. We provide sample videos in assets/examples.
+vipe infer YOUR_VIDEO.mp4
+# Additional options:
+#   --output: Output directory (default: vipe_results)
+#   --visualize: Enable visualization of intermediate and final results (default: false)
+#   --pipeline: Pipeline configuration to use (default: default)
+```
+
+![vipe-vis](assets/vipe-vis.gif)
+
+Currently, we support the following pipeline configurations:
+- `default`: The default pipeline for pinhole cameras.
+- `lyra`: Configuration for results in the [Lyra](https://github.com/nv-tlabs/lyra) paper.
+- 🔥🔥 `dav3`: Using the newest Depth-Anything-V3 model as depth estimation model.
+- `no_vda`: If running video-depth-anything is too memory-consuming for you, this configuration can produce less temporally-stable depth (but empirically more 3D consistent) maps.
+- `wide_angle`: If your video contains some wide-angle or fisheye distortion.
+- `panorama`: For 360° videos, this is only available in the `panorama` branch for now.
+
+One can visualize the results that ViPE produces by running (supported by `viser`):
+```bash
+vipe visualize vipe_results/
+# Please modify the above vipe_results/ path to the output directory of your choice.
+```
+
+![vipe-viser](assets/vipe-viser.gif)
+
+### Using the `run.py` script
+
+The `run.py` script is a more flexible way to run ViPE. Compared to the CLI, the script supports running on multiple videos at once and allows more fine-grained control over the pipeline with `hydra` configs. It also provides an example of using `vipe` as a library in your own project.
+
+Example usages:
+
+```bash
+# Running the full pipeline.
+python run.py pipeline=default streams=raw_mp4_stream streams.base_path=YOUR_VIDEO_OR_DIR_PATH
+
+# Running the pose-only pipeline without depth estimation.
+python run.py pipeline=default streams=raw_mp4_stream streams.base_path=YOUR_VIDEO_OR_DIR_PATH pipeline.post.depth_align_model=null
+```
+
+### Converting to COLMAP format
+
+You can use the following script to convert the ViPE results to COLMAP format. For example:
+```bash
+python scripts/vipe_to_colmap.py vipe_results/ --sequence dog_example
+```
+This will unproject the dense depth maps to create the 3D point cloud. 
+Alternatively for a more lightweight and 3D consistent point cloud, you can add the `--use_slam_map` flag to the above command. This requires you to run the full pipeline with `pipeline.output.save_slam_map=true` to save the additional information.
+
+## Downloading the Dataset
+
+![dataset](assets/dataset.gif)
+
+Together with ViPE we release a large-scale dataset containing ~1M high-quality videos with accurate camera poses and dense depth maps. Specifications of the datasets are listed below:
+
+| Dataset Name   | # Videos | # Frames | Hugging Face Link                                            | License      | Prefix |
+| -------------- | -------- | -------- | ------------------------------------------------------------ | ------------ | ------ |
+| Dynpose-100K++ | 99,501   | 15.8M    | [Link](https://huggingface.co/datasets/nvidia/vipe-dynpose-100kpp) | CC-BY-NC 4.0 | `dpsp` |
+| Wild-SDG-1M    | 966,448  | 78.2M    | [Link](https://huggingface.co/datasets/nvidia/vipe-wild-sdg-1m) | CC-BY-NC 4.0 | `wsdg` |
+| Web360         | 2,114    | 212K     | [Link](https://huggingface.co/datasets/nvidia/vipe-web360)   | CC-BY 4.0    | `w360` |
+
+You can download the datasets using the following utility script:
+
+```bash
+# Replace YOUR_PREFIX with the prefix of the dataset to be downloaded (see prefix column in the table above)
+# You can also use more specific prefixes, e.g. wsdg-003e2c86 to download a specific shard of the dataset.
+python scripts/download_dataset.py --prefix YOUR_PREFIX --output_base YOUR_OUTPUT_DIR --rgb --depth
+```
+
+> Note that the depth component is very large and you might expect a long downloading time. For `rgb` component of the Dynpose-100K++ dataset, we directly retrieve the RGB frames from YouTube. You have to `pip install yt_dlp ffmpeg-python` to use this feature. Please refer to the original [Dynpose-100K dataset](https://huggingface.co/datasets/nvidia/dynpose-100k) for alternative approaches to retrieve the videos.
+
+The dataset itself can be visualized using the same visualization script:
+```bash
+vipe visualize YOUR_OUTPUT_DIR
+```
+
+## Acknowledgments
+
+ViPE is built on top of many great open-source research projects and codebases. Some of these include (not exhaustive):
+- [DROID-SLAM](https://github.com/princeton-vl/DROID-SLAM)
+- [Depth Anything V2](https://github.com/DepthAnything/Depth-Anything-V2)
+- [Metric3Dv2](https://github.com/YvanYin/Metric3D)
+- [PriorDA](https://github.com/SpatialVision/Prior-Depth-Anything)
+- [UniDepth](https://github.com/lpiccinelli-eth/UniDepth)
+- [VideoDepthAnything](https://github.com/DepthAnything/Video-Depth-Anything)
+- [GeoCalib](https://github.com/cvg/GeoCalib)
+- [Segment and Track Anything](https://github.com/z-x-yang/Segment-and-Track-Anything)
+
+Please refer to the [THIRD_PARTY_LICENSES.md](THIRD_PARTY_LICENSES.md) for a full list of projects and their licenses.
+
+We thank useful discussions from Aigul Dzhumamuratova, Viktor Kuznetsov, Soha Pouya, and Ming-Yu Liu, as well as release support from Vishal Kulkarni.
+
+## TODO
+
+- [x] Initial code released under Apache 2.0 license.
+- [x] Full dataset uploaded to Hugging Face for download.
+- [x] Add instructions to run inference on wide-angle and 360° videos.
+- [ ] Add instructions for benchmarking.
+
+## Citation
+
+If you find ViPE useful in your research or application, please consider citing the following whitepaper:
+
+```
+@inproceedings{huang2025vipe,
+    title={ViPE: Video Pose Engine for 3D Geometric Perception},
+    author={Huang, Jiahui and Zhou, Qunjie and Rabeti, Hesam and Korovko, Aleksandr and Ling, Huan and Ren, Xuanchi and Shen, Tianchang and Gao, Jun and Slepichev, Dmitry and Lin, Chen-Hsuan and Ren, Jiawei and Xie, Kevin and Biswas, Joydeep and Leal-Taixe, Laura and Fidler, Sanja},
+    booktitle={NVIDIA Research Whitepapers arXiv:2508.10934},
+    year={2025}
+}
+```
+
+## License
+
+This project will download and install additional third-party **models and softwares**. Note that these models or softwares are not distributed by NVIDIA. Review the license terms of these models and projects before use. This source code is released under the [Apache 2 License](https://www.apache.org/licenses/LICENSE-2.0).

third_party/vipe/configs/default.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - _self_
+  - streams: raw_mp4_stream
+  - pipeline: ???
+  - override hydra/hydra_logging: none
+  - override hydra/job_logging: disabled
+
+ray: false
+prefilter: false
+
+hydra:
+  output_subdir: null
+  job:
+    chdir: false
+  run:
+    dir: .

third_party/vipe/configs/pipeline/dav3.yaml

@@ -0,0 +1,8 @@
+defaults:
+  - default
+
+slam:
+  keyframe_depth: dav3
+
+post:
+  depth_align_model: "mvd_dav3"

third_party/vipe/configs/pipeline/default.yaml

@@ -0,0 +1,47 @@
+defaults:
+  - /slam: default
+
+instance: vipe.pipeline.default.DefaultAnnotationPipeline
+
+# Init configs
+init:
+  camera_type: "pinhole"
+  intrinsics: "geocalib"
+  instance:
+    kf_gap_sec: 2.0
+    phrases: 
+      - person
+      - animal
+      - vehicle
+      - ball
+      - balloon
+      - gun
+      - pet
+      - car
+      - bus
+    add_sky: true
+
+slam:
+  keyframe_depth: unidepth-l
+  optimize_intrinsics: ${neq:${..init.intrinsics},"gt"}
+
+# Post-processing configs
+post:
+  depth_align_model: "adaptive_unidepth-l_svda"
+
+# Output configs
+output:
+  # Path to save results
+  path: vipe_results/
+  skip_exists: false
+
+  # Save artifacts and a info file
+  save_artifacts: false
+
+  # Save reconstruction from SLAM
+  save_slam_map: false
+
+  # Visualization videos to BASE_PATH/vipe/xxx.mp4
+  save_viz: true
+  viz_downsample: 2
+  viz_attributes: [['rgb', 'instance'], ['depth', 'pcd']]

third_party/vipe/configs/pipeline/dynhamr_cameras.yaml

@@ -0,0 +1,9 @@
+# Dyn-HaMR consumes camera poses and intrinsics, not dense depth artifacts.
+defaults:
+  - default
+
+post:
+  depth_align_model: null
+
+output:
+  cameras_only: true

third_party/vipe/configs/pipeline/lyra.yaml

@@ -0,0 +1,8 @@
+defaults:
+  - default
+
+slam:
+  keyframe_depth: moge
+
+post:
+  depth_align_model: "adaptive_moge_vda"

third_party/vipe/configs/pipeline/no_vda.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - default
+
+post:
+  depth_align_model: "adaptive_unidepth-l"

third_party/vipe/configs/pipeline/static_vda.yaml

@@ -0,0 +1,8 @@
+defaults:
+  - default
+
+init:
+  instance: null
+
+post:
+  depth_align_model: "adaptive_unidepth-l_vda"

third_party/vipe/configs/pipeline/wide_angle.yaml

@@ -0,0 +1,14 @@
+defaults:
+  - default
+
+init:
+  camera_type: "mei"
+
+slam:
+  keyframe_depth: unidepth-l
+
+post:
+  depth_align_model: null
+
+output:
+  viz_attributes: [['rgb', 'instance'], ['rectified', 'pcd']]

third_party/vipe/configs/slam/default.yaml

@@ -0,0 +1,53 @@
+buffer: 1024
+# weight for translation / rotation components of flow
+beta: 0.3
+# how much motion before considering new keyframe
+filter_thresh: 2.4
+# number of warmup frames
+warmup: 8
+# threshold to create a new keyframe
+keyframe_thresh: 4.0
+# add edges between frames whithin this distance
+frontend_thresh: 16.0
+# frontend optimization window
+frontend_window: 25
+# force edges between frames within radius
+frontend_radius: 2
+# non-maximal supression of edges
+frontend_nms: 1
+# Whether to initialize sequentially
+seq_init: true
+# Keyframe ids for running backends
+frontend_backend_iters: [16, 64, 256]
+backend_thresh: 22.0
+backend_radius: 2
+backend_nms: 3
+backend_iters: 24
+# Initial values for disparity
+init_disp: 1.0
+# Whether to optimize intrinsics
+optimize_intrinsics: ???
+# Whether to optimize rig
+optimize_rig_rotation: false
+# Whether add reproj factors for cross-views
+cross_view: true
+cross_view_idx: null  # null means default
+adaptive_cross_view: false  # Backend only, recompute cross_view_idx
+# Chunk size of the infill
+infill_chunk_size: 16
+# Whether or not to infill dense disparity
+infill_dense_disp: False
+# Filter threshold for SLAM map (percentage of the mean depth difference allowed)
+# (This is also used for dense disp extraction)
+map_filter_thresh: 0.05
+# Whether to visualize the results using rerun
+visualize: false
+# Use keyframe depth for recovering metric depth
+keyframe_depth: metric3d-small
+
+ba:
+  dense_disp_alpha: 0.001
+  intrinsics_damping_scale: 1.0
+
+sparse_tracks:
+  name: dummy

third_party/vipe/configs/streams/frame_dir_stream.yaml

@@ -0,0 +1,7 @@
+instance: vipe.streams.frame_dir_stream.FrameDirStreamList
+
+base_path: ??
+frame_start: 0
+frame_end: -1
+frame_skip: 1
+cached: false

third_party/vipe/configs/streams/raw_mp4_stream.yaml

@@ -0,0 +1,7 @@
+instance: vipe.streams.raw_mp4_stream.RawMP4StreamList
+
+base_path: ??
+frame_start: 0
+frame_end: 1000
+frame_skip: 1
+cached: false

third_party/vipe/csrc/bind.cpp

@@ -0,0 +1,49 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/extension.h>
+
+void pybind_droid_net_ext(py::module &m);
+void pybind_grounding_dino_ext(py::module &m);
+void pybind_utils_ext(py::module &m);
+void pybind_slam_ext(py::module &m);
+void pybind_scatter_ext(py::module &m);
+void pybind_lietorch_ext(py::module &m);
+void pybind_corr_ext(py::module &m);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    py::module m_droid_net = m.def_submodule("droid_net_ext");
+    pybind_droid_net_ext(m_droid_net);
+
+    py::module m_gdino = m.def_submodule("grounding_dino_ext");
+    pybind_grounding_dino_ext(m_gdino);
+
+    py::module m_utils = m.def_submodule("utils_ext");
+    pybind_utils_ext(m_utils);
+
+    py::module m_slam = m.def_submodule("slam_ext");
+    pybind_slam_ext(m_slam);
+
+    py::module m_scatter = m.def_submodule("scatter_ext");
+    pybind_scatter_ext(m_scatter);
+
+    py::module m_lietorch = m.def_submodule("lietorch_ext");
+    pybind_lietorch_ext(m_lietorch);
+
+    py::module m_corr = m.def_submodule("corr_ext");
+    pybind_corr_ext(m_corr);
+}

third_party/vipe/csrc/corr_ext/correlation.cpp

@@ -0,0 +1,148 @@
+/**
+ * This file includes code originally from the Pytorch correlation repository:
+ * https://github.com/ClementPinard/Pytorch-Correlation-extension、
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+using namespace torch;
+
+#include <vector>
+
+#define WITHIN_BOUNDS(x, y, H, W) (x >= 0 && x < H && y >= 0 && y < W)
+
+template <typename scalar_t>
+static void correlate_patch(TensorAccessor<scalar_t, 3> input1, TensorAccessor<scalar_t, 3> input2, scalar_t *dst,
+                            int kH, int kW, int dilationH, int dilationW, int u, int v, int shiftU, int shiftV) {
+    const int C = input1.size(0);
+    const int iH = input1.size(1);
+    const int iW = input1.size(2);
+    for (int c = 0; c < C; ++c) {
+        for (int i = 0; i < kH; ++i) {
+            int i1 = u + i * dilationH;
+            int i2 = i1 + shiftU;
+            if WITHIN_BOUNDS (i1, i2, iH, iH) {
+                for (int j = 0; j < kW; ++j) {
+                    int j1 = v + j * dilationW;
+                    int j2 = j1 + shiftV;
+                    if WITHIN_BOUNDS (j1, j2, iW, iW) {
+                        scalar_t v1 = input1[c][i1][j1];
+                        scalar_t v2 = input2[c][i2][j2];
+                        *dst += v1 * v2;
+                    }
+                }
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+static void correlate_patch_grad(TensorAccessor<scalar_t, 3> input1, TensorAccessor<scalar_t, 3> gradInput1,
+                                 TensorAccessor<scalar_t, 3> input2, TensorAccessor<scalar_t, 3> gradInput2,
+                                 scalar_t gradOutput, int kH, int kW, int dilationH, int dilationW, int u, int v,
+                                 int shiftU, int shiftV) {
+    const int C = input1.size(0);
+    const int iH = input1.size(1);
+    const int iW = input1.size(2);
+
+    for (int c = 0; c < C; ++c) {
+        for (int i = 0; i < kH; ++i) {
+            int i1 = u + i * dilationH;
+            int i2 = i1 + shiftU;
+            if WITHIN_BOUNDS (i1, i2, iH, iH) {
+                for (int j = 0; j < kW; ++j) {
+                    int j1 = v + j * dilationW;
+                    int j2 = j1 + shiftV;
+                    if WITHIN_BOUNDS (j1, j2, iW, iW) {
+                        scalar_t v1 = input1[c][i1][j1];
+                        scalar_t v2 = input2[c][i2][j2];
+                        gradInput2[c][i2][j2] += gradOutput * v1;
+                        gradInput1[c][i1][j1] += gradOutput * v2;
+                    }
+                }
+            }
+        }
+    }
+}
+
+torch::Tensor correlation_cpp_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
+                                      int patchW, int padH, int padW, int dilationH, int dilationW, int dilation_patchH,
+                                      int dilation_patchW, int dH, int dW) {
+    const auto batch_size = input1.size(0);
+    const auto iH = input1.size(2);
+    const auto iW = input1.size(3);
+    const int patchRadH = (patchH - 1) / 2;
+    const int patchRadW = (patchW - 1) / 2;
+    const int dilatedKH = (kH - 1) * dilationH + 1;
+    const int dilatedKW = (kW - 1) * dilationW + 1;
+
+    const auto oH = (iH + 2 * padH - dilatedKH) / dH + 1;
+    const auto oW = (iW + 2 * padW - dilatedKW) / dW + 1;
+    auto output = at::zeros({batch_size, patchH, patchW, oH, oW}, input1.options());
+
+    int n, ph, pw, h, w;
+#pragma omp parallel for private(n, ph, pw, h, w) collapse(2)
+    for (n = 0; n < batch_size; ++n) {
+        for (ph = 0; ph < patchH; ++ph) {
+            for (pw = 0; pw < patchW; ++pw) {
+                AT_DISPATCH_FLOATING_TYPES(
+                    input1.scalar_type(), "correlation_forward_cpp", ([&] {
+                        auto input1_acc = input1.accessor<scalar_t, 4>();
+                        auto input2_acc = input2.accessor<scalar_t, 4>();
+                        auto output_acc = output.accessor<scalar_t, 5>();
+                        for (h = 0; h < oH; ++h) {
+                            for (w = 0; w < oW; ++w) {
+                                correlate_patch(input1_acc[n], input2_acc[n], &output_acc[n][ph][pw][h][w], kH, kW,
+                                                dilationH, dilationW, -padH + h * dH, -padW + w * dW,
+                                                (ph - patchRadH) * dilation_patchH, (pw - patchRadW) * dilation_patchW);
+                            }
+                        }
+                    }));
+            }
+        }
+    }
+    return output;
+}
+
+std::vector<torch::Tensor> correlation_cpp_backward(torch::Tensor input1, torch::Tensor input2,
+                                                    torch::Tensor gradOutput, int kH, int kW, int patchH, int patchW,
+                                                    int padH, int padW, int dilationH, int dilationW,
+                                                    int dilation_patchH, int dilation_patchW, int dH, int dW) {
+    const int batch_size = input1.size(0);
+    const int patchRadH = (patchH - 1) / 2;
+    const int patchRadW = (patchW - 1) / 2;
+    const int oH = gradOutput.size(3);
+    const int oW = gradOutput.size(4);
+
+    auto gradInput1 = torch::zeros_like(input1);
+
+    auto gradInput2 = torch::zeros_like(input2);
+
+    int n, ph, pw, h, w;
+#pragma omp parallel for private(n, ph, pw, h, w)
+    for (n = 0; n < batch_size; ++n) {
+        AT_DISPATCH_FLOATING_TYPES(
+            input1.scalar_type(), "correlation_backward_cpp", ([&] {
+                auto input1_acc = input1.accessor<scalar_t, 4>();
+                auto gradInput1_acc = gradInput1.accessor<scalar_t, 4>();
+                auto input2_acc = input2.accessor<scalar_t, 4>();
+                auto gradInput2_acc = gradInput2.accessor<scalar_t, 4>();
+                auto gradOutput_acc = gradOutput.accessor<scalar_t, 5>();
+
+                for (ph = 0; ph < patchH; ++ph) {
+                    for (pw = 0; pw < patchW; ++pw) {
+                        for (h = 0; h < oH; ++h) {
+                            for (w = 0; w < oW; ++w) {
+                                correlate_patch_grad(input1_acc[n], gradInput1_acc[n], input2_acc[n], gradInput2_acc[n],
+                                                     gradOutput_acc[n][ph][pw][h][w], kH, kW, dilationH, dilationW,
+                                                     -padH + h * dH, -padW + w * dW, (ph - patchRadH) * dilation_patchH,
+                                                     (pw - patchRadW) * dilation_patchW);
+                            }
+                        }
+                    }
+                }
+            }));
+    }
+
+    return {gradInput1, gradInput2};
+}

third_party/vipe/csrc/corr_ext/correlation_cuda_kernel.cu

@@ -0,0 +1,286 @@
+/**
+ * This file includes code originally from the Pytorch correlation repository:
+ * https://github.com/ClementPinard/Pytorch-Correlation-extension、
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/types.h>
+using namespace torch;
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <vector>
+#include <iostream>
+
+// Cuda tensor accessor definitions
+// restrict pointer traits piroritize speed over memory consumption
+#define TensorAcc4R PackedTensorAccessor32<scalar_t, 4, RestrictPtrTraits>
+#define TensorAcc5R PackedTensorAccessor32<scalar_t, 5, RestrictPtrTraits>
+#define WITHIN_BOUNDS(x, y, H, W) (x >= 0 && x < H && y >= 0 && y < W)
+
+#define THREADS_FORWARD 32
+#define THREADS_BACKWARD 5
+
+namespace corr {
+template <typename scalar_t>
+__global__ void correlation_cuda_forward_kernel(const TensorAcc4R rInput1, const TensorAcc4R rInput2,
+                                                TensorAcc5R output, int kH, int kW, int patchH, int patchW, int padH,
+                                                int padW, int dilationH, int dilationW, int dilation_patchH,
+                                                int dilation_patchW, int dH, int dW) {
+    const int iH = rInput1.size(1);
+    const int iW = rInput1.size(2);
+    const int C = rInput1.size(3);
+
+    const int n = blockIdx.x;
+    const int h = blockIdx.y;
+    const int w = blockIdx.z;
+    const int thread = threadIdx.x;
+
+    const int start_i = -padH + h * dH;
+    const int start_j = -padW + w * dW;
+
+    const int patchRadH = dilation_patchH * (patchH - 1) / 2;
+    const int patchRadW = dilation_patchW * (patchW - 1) / 2;
+
+    __shared__ scalar_t prod_sum[THREADS_FORWARD];
+
+    for (int ph = 0; ph < patchH; ++ph) {
+        int ph_dilated = ph * dilation_patchH - patchRadH;
+        for (int pw = 0; pw < patchW; ++pw) {
+            int pw_dilated = pw * dilation_patchW - patchRadW;
+            prod_sum[thread] = 0;
+            for (int i = 0; i < kH; ++i) {
+                int i1 = start_i + i * dilationH;
+                int i2 = i1 + ph_dilated;
+                if WITHIN_BOUNDS (i1, i2, iH, iH) {
+                    for (int j = 0; j < kW; ++j) {
+                        int j1 = start_j + j * dilationW;
+                        int j2 = j1 + pw_dilated;
+                        if WITHIN_BOUNDS (j1, j2, iW, iW) {
+                            for (int c = thread; c < C; c += THREADS_FORWARD) {
+                                scalar_t v1 = rInput1[n][i1][j1][c];
+                                scalar_t v2 = rInput2[n][i2][j2][c];
+                                prod_sum[thread] += v1 * v2;
+                            }
+                        }
+                    }
+                }
+            }
+            // accumulate
+            __syncthreads();
+            if (thread == 0) {
+                scalar_t reduce_sum = 0;
+                for (int index = 0; index < THREADS_FORWARD; ++index) {
+                    reduce_sum += prod_sum[index];
+                }
+                output[n][ph][pw][h][w] = reduce_sum;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void correlation_cuda_backward_kernel_input1(const TensorAcc5R gradOutput, const TensorAcc4R input2,
+                                                        TensorAcc4R gradInput1, const int kH, const int kW,
+                                                        const int patchH, const int patchW, const int padH,
+                                                        const int padW, const int dilationH, const int dilationW,
+                                                        const int dilation_patchH, const int dilation_patchW,
+                                                        const int dH, const int dW, const int batch) {
+    const int iH = input2.size(2);
+    const int iW = input2.size(3);
+
+    const int H = gradOutput.size(3);
+    const int W = gradOutput.size(4);
+
+    const int patchRadH = (patchH - 1) / 2;
+    const int patchRadW = (patchW - 1) / 2;
+
+    const int n = batch;
+    const int c = blockIdx.x;
+    const int h = blockIdx.y;
+    const int w = blockIdx.z;
+    const int ph_off = threadIdx.x;
+    const int pw_off = threadIdx.y;
+
+    const int h_2 = h + padH;
+    const int w_2 = w + padW;
+    const int min_h = h_2 - kH * dilationH;
+    const int min_w = w_2 - kW * dilationW;
+
+    __shared__ scalar_t prod_sum[THREADS_BACKWARD][THREADS_BACKWARD];
+    prod_sum[ph_off][pw_off] = 0;
+
+    for (int ph = ph_off; ph < patchH; ph += THREADS_BACKWARD) {
+        int i1 = h + dilation_patchH * (ph - patchRadH);
+        for (int pw = pw_off; pw < patchW; pw += THREADS_BACKWARD) {
+            int j1 = w + dilation_patchW * (pw - patchRadW);
+            if (WITHIN_BOUNDS(i1, j1, iH, iW)) {
+                scalar_t val = input2[n][c][i1][j1];
+                for (int h_3 = h_2; h_3 > min_h; h_3 -= dilationH) {
+                    int i2 = (h_3) / dH;
+                    if (i2 * dH != h_3) continue;
+                    for (int w_3 = w_2; w_3 > min_w; w_3 -= dilationW) {
+                        int j2 = (w_3) / dW;
+                        if (j2 * dW != w_3) continue;
+                        if WITHIN_BOUNDS (i2, j2, H, W) {
+                            prod_sum[ph_off][pw_off] += gradOutput[n][ph][pw][i2][j2] * val;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    __syncthreads();
+
+    if (ph_off == 0 && pw_off == 0) {
+        scalar_t reduce_sum = 0;
+        for (int ph = 0; ph < THREADS_BACKWARD; ++ph) {
+            for (int pw = 0; pw < THREADS_BACKWARD; ++pw) {
+                reduce_sum += prod_sum[ph][pw];
+            }
+        }
+        gradInput1[n][c][h][w] = reduce_sum;
+    }
+}
+
+template <typename scalar_t>
+__global__ void correlation_cuda_backward_kernel_input2(const TensorAcc5R gradOutput, const TensorAcc4R input1,
+                                                        TensorAcc4R gradInput2, int kH, int kW, int patchH, int patchW,
+                                                        int padH, int padW, int dilationH, int dilationW,
+                                                        int dilation_patchH, int dilation_patchW, int dH, int dW,
+                                                        int batch) {
+    const int iH = input1.size(2);
+    const int iW = input1.size(3);
+
+    const int patchRadH = (patchH - 1) / 2;
+    const int patchRadW = (patchW - 1) / 2;
+
+    const int H = gradOutput.size(3);
+    const int W = gradOutput.size(4);
+
+    const int dilatedKH = kH * dilationH;
+    const int dilatedKW = kW * dilationW;
+
+    const int n = batch;
+    const int c = blockIdx.x;
+    const int h = blockIdx.y;
+    const int w = blockIdx.z;
+    const int ph_off = threadIdx.x;
+    const int pw_off = threadIdx.y;
+
+    __shared__ scalar_t prod_sum[THREADS_BACKWARD][THREADS_BACKWARD];
+    prod_sum[ph_off][pw_off] = 0;
+
+    for (int ph = ph_off; ph < patchH; ph += THREADS_BACKWARD) {
+        int i1 = h - dilation_patchH * (ph - patchRadH);
+        for (int pw = pw_off; pw < patchW; pw += THREADS_BACKWARD) {
+            int j1 = w - dilation_patchW * (pw - patchRadW);
+            if WITHIN_BOUNDS (i1, j1, iH, iW) {
+                scalar_t val = input1[n][c][i1][j1];
+
+                const int h_2 = i1 + padH;
+                const int w_2 = j1 + padW;
+                const int min_h = h_2 - dilatedKH;
+                const int min_w = w_2 - dilatedKW;
+
+                for (int h_3 = h_2; h_3 > min_h; h_3 -= dilationH) {
+                    int i2 = (h_3) / dH;
+                    if (i2 * dH != h_3) continue;
+                    for (int w_3 = w_2; w_3 > min_w; w_3 -= dilationW) {
+                        int j2 = (w_3) / dW;
+                        if (j2 * dW != w_3) continue;
+                        if WITHIN_BOUNDS (i2, j2, H, W) {
+                            prod_sum[ph_off][pw_off] += gradOutput[n][ph][pw][i2][j2] * val;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    __syncthreads();
+
+    if (ph_off == 0 && pw_off == 0) {
+        scalar_t reduce_sum = 0;
+        for (int ph = 0; ph < THREADS_BACKWARD; ++ph) {
+            for (int pw = 0; pw < THREADS_BACKWARD; ++pw) {
+                reduce_sum += prod_sum[ph][pw];
+            }
+        }
+        gradInput2[n][c][h][w] = reduce_sum;
+    }
+}
+}  // namespace corr
+
+torch::Tensor correlation_cuda_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
+                                       int patchW, int padH, int padW, int dilationH, int dilationW,
+                                       int dilation_patchH, int dilation_patchW, int dH, int dW) {
+    const int batch_size = input1.size(0);
+    const int iH = input1.size(2);
+    const int iW = input1.size(3);
+    const int dilatedKH = (kH - 1) * dilationH + 1;
+    const int dilatedKW = (kW - 1) * dilationW + 1;
+
+    const auto oH = (iH + 2 * padH - dilatedKH) / dH + 1;
+    const auto oW = (iW + 2 * padW - dilatedKW) / dW + 1;
+    auto output = torch::zeros({batch_size, patchH, patchW, oH, oW}, input1.options());
+
+    auto trInput1 = input1.permute({0, 2, 3, 1}).contiguous();
+    auto trInput2 = input2.permute({0, 2, 3, 1}).contiguous();
+
+    const int threads = THREADS_FORWARD;
+    const dim3 blocks(batch_size, oH, oW);
+
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half, at::ScalarType::BFloat16, input1.scalar_type(), "correlation_forward_cuda", ([&] {
+            TensorAcc4R trInput1_acc = trInput1.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc4R trInput2_acc = trInput2.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc5R output_acc = output.packed_accessor32<scalar_t, 5, RestrictPtrTraits>();
+            corr::correlation_cuda_forward_kernel<scalar_t>
+                <<<blocks, threads>>>(trInput1_acc, trInput2_acc, output_acc, kH, kW, patchH, patchW, padH, padW,
+                                      dilationH, dilationW, dilation_patchH, dilation_patchW, dH, dW);
+        }));
+
+    return output;
+}
+
+std::vector<torch::Tensor> correlation_cuda_backward(torch::Tensor input1, torch::Tensor input2,
+                                                     torch::Tensor gradOutput, int kH, int kW, int patchH, int patchW,
+                                                     int padH, int padW, int dilationH, int dilationW,
+                                                     int dilation_patchH, int dilation_patchW, int dH, int dW) {
+    auto gradInput1 = torch::zeros_like(input1);
+    auto gradInput2 = torch::zeros_like(input2);
+
+    const int batch_size = input1.size(0);
+    const int iH = input1.size(2);
+    const int iW = input1.size(3);
+    const int C = input1.size(1);
+
+    const dim3 blocks(C, iH, iW);
+    const dim3 threads(THREADS_BACKWARD, THREADS_BACKWARD);
+
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half, at::ScalarType::BFloat16, input1.scalar_type(), "correlation_backward_cuda", ([&] {
+            TensorAcc4R input1_acc = input1.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc4R input2_acc = input2.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc4R gradInput1_acc = gradInput1.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc4R gradInput2_acc = gradInput2.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc5R gradOutput_acc = gradOutput.packed_accessor32<scalar_t, 5, RestrictPtrTraits>();
+
+            for (int n = 0; n < batch_size; ++n) {
+                corr::correlation_cuda_backward_kernel_input1<scalar_t>
+                    <<<blocks, threads>>>(gradOutput_acc, input2_acc, gradInput1_acc, kH, kW, patchH, patchW, padH,
+                                          padW, dilationH, dilationW, dilation_patchH, dilation_patchW, dH, dW, n);
+            }
+
+            for (int n = 0; n < batch_size; ++n) {
+                corr::correlation_cuda_backward_kernel_input2<scalar_t>
+                    <<<blocks, threads>>>(gradOutput_acc, input1_acc, gradInput2_acc, kH, kW, patchH, patchW, padH,
+                                          padW, dilationH, dilationW, dilation_patchH, dilation_patchW, dH, dW, n);
+            }
+        }));
+
+    return {gradInput1, gradInput2};
+}

third_party/vipe/csrc/corr_ext/correlation_sampler.cpp

@@ -0,0 +1,85 @@
+/**
+ * This file includes code originally from the Pytorch correlation repository:
+ * https://github.com/ClementPinard/Pytorch-Correlation-extension、
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <vector>
+#include <iostream>
+
+// declarations
+
+torch::Tensor correlation_cpp_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
+                                      int patchW, int padH, int padW, int dilationH, int dilationW, int dilation_patchH,
+                                      int dilation_patchW, int dH, int dW);
+
+std::vector<torch::Tensor> correlation_cpp_backward(torch::Tensor grad_output, torch::Tensor input1,
+                                                    torch::Tensor input2, int kH, int kW, int patchH, int patchW,
+                                                    int padH, int padW, int dilationH, int dilationW,
+                                                    int dilation_patchH, int dilation_patchW, int dH, int dW);
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x, " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous")
+#define CHECK_INPUT(x) \
+    CHECK_CUDA(x);     \
+    CHECK_CONTIGUOUS(x)
+#define CHECK_SAME_DEVICE(x, y) TORCH_CHECK(x.device() == y.device(), #x " is not on same device as " #y)
+
+torch::Tensor correlation_cuda_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
+                                       int patchW, int padH, int padW, int dilationH, int dilationW,
+                                       int dilation_patchH, int dilation_patchW, int dH, int dW);
+
+std::vector<torch::Tensor> correlation_cuda_backward(torch::Tensor grad_output, torch::Tensor input1,
+                                                     torch::Tensor input2, int kH, int kW, int patchH, int patchW,
+                                                     int padH, int padW, int dilationH, int dilationW,
+                                                     int dilation_patchH, int dilation_patchW, int dH, int dW);
+
+// C++ interface
+
+torch::Tensor correlation_sample_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
+                                         int patchW, int padH, int padW, int dilationH, int dilationW,
+                                         int dilation_patchH, int dilation_patchW, int dH, int dW) {
+    if (input1.device().is_cuda()) {
+        CHECK_INPUT(input1);
+        CHECK_INPUT(input2);
+
+        // set device of input1 as default CUDA device
+        // https://pytorch.org/cppdocs/api/structc10_1_1cuda_1_1_optional_c_u_d_a_guard.html
+        const at::cuda::OptionalCUDAGuard guard_input1(device_of(input1));
+        CHECK_SAME_DEVICE(input1, input2);
+
+        return correlation_cuda_forward(input1, input2, kH, kW, patchH, patchW, padH, padW, dilationH, dilationW,
+                                        dilation_patchH, dilation_patchW, dH, dW);
+    } else {
+        return correlation_cpp_forward(input1, input2, kH, kW, patchH, patchW, padH, padW, dilationH, dilationW,
+                                       dilation_patchH, dilation_patchW, dH, dW);
+    }
+}
+
+std::vector<torch::Tensor> correlation_sample_backward(torch::Tensor input1, torch::Tensor input2,
+                                                       torch::Tensor grad_output, int kH, int kW, int patchH,
+                                                       int patchW, int padH, int padW, int dilationH, int dilationW,
+                                                       int dilation_patchH, int dilation_patchW, int dH, int dW) {
+    if (grad_output.device().is_cuda()) {
+        CHECK_INPUT(input1);
+        CHECK_INPUT(input2);
+
+        // set device of input1 as default CUDA device
+        const at::cuda::OptionalCUDAGuard guard_input1(device_of(input1));
+        CHECK_SAME_DEVICE(input1, input2);
+        CHECK_SAME_DEVICE(input1, grad_output);
+
+        return correlation_cuda_backward(input1, input2, grad_output, kH, kW, patchH, patchW, padH, padW, dilationH,
+                                         dilationW, dilation_patchH, dilation_patchW, dH, dW);
+    } else {
+        return correlation_cpp_backward(input1, input2, grad_output, kH, kW, patchH, patchW, padH, padW, dilationH,
+                                        dilationW, dilation_patchH, dilation_patchW, dH, dW);
+    }
+}
+
+void pybind_corr_ext(py::module& m) {
+    m.def("forward", &correlation_sample_forward, "Spatial Correlation Sampler Forward");
+    m.def("backward", &correlation_sample_backward, "Spatial Correlation Sampler backward");
+}

third_party/vipe/csrc/droid_net_ext/altcorr_kernel.cu

@@ -0,0 +1,321 @@
+/**
+ * This file includes code originally from the DROID-SLAM repository:
+ * https://github.com/princeton-vl/DROID-SLAM
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <ATen/ATen.h>
+#include <ATen/NativeFunctions.h>
+#include <ATen/cuda/CUDAApplyUtils.cuh>
+#include <ATen/native/cuda/KernelUtils.cuh>
+
+#define BLOCK_H 4
+#define BLOCK_W 8
+#define BLOCK_HW BLOCK_H *BLOCK_W
+#define CHANNEL_STRIDE 32
+
+__forceinline__ __device__ bool within_bounds(int h, int w, int H, int W) { return h >= 0 && h < H && w >= 0 && w < W; }
+
+template <typename scalar_t>
+__global__ void altcorr_forward_kernel(const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap1,
+                                       const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap2,
+                                       const torch::PackedTensorAccessor32<float, 5, torch::RestrictPtrTraits> coords,
+                                       torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr,
+                                       int r) {
+    const int b = blockIdx.x;
+    const int h0 = blockIdx.y * blockDim.x;
+    const int w0 = blockIdx.z * blockDim.y;
+    const int tid = threadIdx.x * blockDim.y + threadIdx.y;
+
+    const int H1 = fmap1.size(1);
+    const int W1 = fmap1.size(2);
+    const int H2 = fmap2.size(1);
+    const int W2 = fmap2.size(2);
+    const int N = coords.size(1);
+    const int C = fmap1.size(3);
+
+    __shared__ scalar_t f1[CHANNEL_STRIDE][BLOCK_HW];
+    __shared__ scalar_t f2[CHANNEL_STRIDE][BLOCK_HW];
+
+    __shared__ float x2s[BLOCK_HW];
+    __shared__ float y2s[BLOCK_HW];
+
+    for (int c = 0; c < C; c += CHANNEL_STRIDE) {
+        for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+            int k1 = k + tid / CHANNEL_STRIDE;
+            int h1 = h0 + k1 / BLOCK_W;
+            int w1 = w0 + k1 % BLOCK_W;
+            int c1 = tid % CHANNEL_STRIDE;
+
+            if (within_bounds(h1, w1, H1, W1))
+                f1[c1][k1] = fmap1[b][h1][w1][c + c1];
+
+            else
+                f1[c1][k1] = 0.0;
+        }
+
+        __syncthreads();
+
+        for (int n = 0; n < N; n++) {
+            int h1 = h0 + threadIdx.x;
+            int w1 = w0 + threadIdx.y;
+            if (within_bounds(h1, w1, H1, W1)) {
+                x2s[tid] = coords[b][n][h1][w1][0];
+                y2s[tid] = coords[b][n][h1][w1][1];
+            }
+
+            float dx = x2s[tid] - floor(x2s[tid]);
+            float dy = y2s[tid] - floor(y2s[tid]);
+
+            int rd = 2 * r + 1;
+            for (int iy = 0; iy < rd + 1; iy++) {
+                for (int ix = 0; ix < rd + 1; ix++) {
+                    for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+                        int k1 = k + tid / CHANNEL_STRIDE;
+                        int h2 = static_cast<int>(floor(y2s[k1])) - r + iy;
+                        int w2 = static_cast<int>(floor(x2s[k1])) - r + ix;
+                        int c2 = tid % CHANNEL_STRIDE;
+
+                        if (within_bounds(h2, w2, H2, W2))
+                            f2[c2][k1] = fmap2[b][h2][w2][c + c2];
+
+                        else
+                            f2[c2][k1] = static_cast<scalar_t>(0.0);
+                    }
+
+                    __syncthreads();
+
+                    scalar_t s = 0.0;
+                    for (int k = 0; k < CHANNEL_STRIDE; k++) s += f1[k][tid] * f2[k][tid];
+
+                    int ix_nw = H1 * W1 * ((iy - 1) + rd * (ix - 1));
+                    int ix_ne = H1 * W1 * ((iy - 1) + rd * ix);
+                    int ix_sw = H1 * W1 * (iy + rd * (ix - 1));
+                    int ix_se = H1 * W1 * (iy + rd * ix);
+
+                    // int ix_nw = ((iy-1) + rd*(ix-1));
+                    // int ix_ne = ((iy-1) + rd*ix);
+                    // int ix_sw = (iy + rd*(ix-1));
+                    // int ix_se = (iy + rd*ix);
+
+                    scalar_t nw = s * static_cast<scalar_t>((dy) * (dx));
+                    scalar_t ne = s * static_cast<scalar_t>((dy) * (1 - dx));
+                    scalar_t sw = s * static_cast<scalar_t>((1 - dy) * (dx));
+                    scalar_t se = s * static_cast<scalar_t>((1 - dy) * (1 - dx));
+
+                    // if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1))
+                    //   corr[b][n][ix_nw][h1][w1] += nw;
+
+                    // if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1))
+                    //   corr[b][n][ix_ne][h1][w1] += ne;
+
+                    // if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1))
+                    //   corr[b][n][ix_sw][h1][w1] += sw;
+
+                    // if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1))
+                    //   corr[b][n][ix_se][h1][w1] += se;
+
+                    scalar_t *corr_ptr = &corr[b][n][0][h1][w1];
+
+                    if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_nw) += nw;
+
+                    if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_ne) += ne;
+
+                    if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_sw) += sw;
+
+                    if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_se) += se;
+                }
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void altcorr_backward_kernel(
+    const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap1,
+    const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap2,
+    const torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> coords,
+    const torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr_grad,
+    torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap1_grad,
+    torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap2_grad,
+    torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> coords_grad, int r) {
+    const int b = blockIdx.x;
+    const int h0 = blockIdx.y * blockDim.x;
+    const int w0 = blockIdx.z * blockDim.y;
+    const int tid = threadIdx.x * blockDim.y + threadIdx.y;
+
+    const int H1 = fmap1.size(1);
+    const int W1 = fmap1.size(2);
+    const int H2 = fmap2.size(1);
+    const int W2 = fmap2.size(2);
+    const int N = coords.size(1);
+    const int C = fmap1.size(3);
+
+    __shared__ scalar_t f1[CHANNEL_STRIDE][BLOCK_HW + 1];
+    __shared__ scalar_t f2[CHANNEL_STRIDE][BLOCK_HW + 1];
+
+    __shared__ scalar_t f1_grad[CHANNEL_STRIDE][BLOCK_HW + 1];
+    __shared__ scalar_t f2_grad[CHANNEL_STRIDE][BLOCK_HW + 1];
+
+    __shared__ scalar_t x2s[BLOCK_HW];
+    __shared__ scalar_t y2s[BLOCK_HW];
+
+    for (int c = 0; c < C; c += CHANNEL_STRIDE) {
+        for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+            int k1 = k + tid / CHANNEL_STRIDE;
+            int h1 = h0 + k1 / BLOCK_W;
+            int w1 = w0 + k1 % BLOCK_W;
+            int c1 = tid % CHANNEL_STRIDE;
+
+            auto fptr = fmap1[b][h1][w1];
+            if (within_bounds(h1, w1, H1, W1))
+                f1[c1][k1] = fptr[c + c1];
+            else
+                f1[c1][k1] = 0.0;
+
+            f1_grad[c1][k1] = 0.0;
+        }
+
+        __syncthreads();
+
+        int h1 = h0 + threadIdx.x;
+        int w1 = w0 + threadIdx.y;
+
+        for (int n = 0; n < N; n++) {
+            x2s[tid] = coords[b][n][h1][w1][0];
+            y2s[tid] = coords[b][n][h1][w1][1];
+
+            scalar_t dx = x2s[tid] - floor(x2s[tid]);
+            scalar_t dy = y2s[tid] - floor(y2s[tid]);
+
+            int rd = 2 * r + 1;
+            for (int iy = 0; iy < rd + 1; iy++) {
+                for (int ix = 0; ix < rd + 1; ix++) {
+                    for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+                        int k1 = k + tid / CHANNEL_STRIDE;
+                        int h2 = static_cast<int>(floor(y2s[k1])) - r + iy;
+                        int w2 = static_cast<int>(floor(x2s[k1])) - r + ix;
+                        int c2 = tid % CHANNEL_STRIDE;
+
+                        auto fptr = fmap2[b][h2][w2];
+                        if (within_bounds(h2, w2, H2, W2))
+                            f2[c2][k1] = fptr[c + c2];
+                        else
+                            f2[c2][k1] = 0.0;
+
+                        f2_grad[c2][k1] = 0.0;
+                    }
+
+                    __syncthreads();
+
+                    const scalar_t *grad_ptr = &corr_grad[b][n][0][h1][w1];
+                    scalar_t g = 0.0;
+
+                    int ix_nw = H1 * W1 * ((iy - 1) + rd * (ix - 1));
+                    int ix_ne = H1 * W1 * ((iy - 1) + rd * ix);
+                    int ix_sw = H1 * W1 * (iy + rd * (ix - 1));
+                    int ix_se = H1 * W1 * (iy + rd * ix);
+
+                    if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1)) g += *(grad_ptr + ix_nw) * dy * dx;
+
+                    if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1)) g += *(grad_ptr + ix_ne) * dy * (1 - dx);
+
+                    if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1)) g += *(grad_ptr + ix_sw) * (1 - dy) * dx;
+
+                    if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1))
+                        g += *(grad_ptr + ix_se) * (1 - dy) * (1 - dx);
+
+                    for (int k = 0; k < CHANNEL_STRIDE; k++) {
+                        f1_grad[k][tid] += g * f2[k][tid];
+                        f2_grad[k][tid] += g * f1[k][tid];
+                    }
+
+                    for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+                        int k1 = k + tid / CHANNEL_STRIDE;
+                        int h2 = static_cast<int>(floor(y2s[k1])) - r + iy;
+                        int w2 = static_cast<int>(floor(x2s[k1])) - r + ix;
+                        int c2 = tid % CHANNEL_STRIDE;
+
+                        scalar_t *fptr = &fmap2_grad[b][h2][w2][0];
+                        if (within_bounds(h2, w2, H2, W2)) atomicAdd(fptr + c + c2, f2_grad[c2][k1]);
+                    }
+                }
+            }
+        }
+        __syncthreads();
+
+        for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+            int k1 = k + tid / CHANNEL_STRIDE;
+            int h1 = h0 + k1 / BLOCK_W;
+            int w1 = w0 + k1 % BLOCK_W;
+            int c1 = tid % CHANNEL_STRIDE;
+
+            scalar_t *fptr = &fmap1_grad[b][h1][w1][0];
+            if (within_bounds(h1, w1, H1, W1)) fptr[c + c1] += f1_grad[c1][k1];
+        }
+    }
+}
+
+std::vector<torch::Tensor> altcorr_cuda_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
+                                                int radius) {
+    const auto B = coords.size(0);
+    const auto N = coords.size(1);
+    const auto H = coords.size(2);
+    const auto W = coords.size(3);
+
+    const auto rd = 2 * radius + 1;
+    auto opts = fmap1.options();
+    auto corr = torch::zeros({B, N, rd * rd, H, W}, opts);
+
+    const dim3 blocks(B, (H + BLOCK_H - 1) / BLOCK_H, (W + BLOCK_W - 1) / BLOCK_W);
+    const dim3 threads(BLOCK_H, BLOCK_W);
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(fmap1.scalar_type(), "altcorr_forward_kernel", ([&] {
+                                            altcorr_forward_kernel<scalar_t><<<blocks, threads>>>(
+                                                fmap1.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+                                                fmap2.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+                                                coords.packed_accessor32<float, 5, torch::RestrictPtrTraits>(),
+                                                corr.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
+                                                radius);
+                                        }));
+
+    return {corr};
+}
+
+std::vector<torch::Tensor> altcorr_cuda_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
+                                                 torch::Tensor corr_grad, int radius) {
+    const auto B = coords.size(0);
+    const auto N = coords.size(1);
+
+    const auto H1 = fmap1.size(1);
+    const auto W1 = fmap1.size(2);
+    const auto H2 = fmap2.size(1);
+    const auto W2 = fmap2.size(2);
+    const auto C = fmap1.size(3);
+
+    auto opts = fmap1.options();
+    auto fmap1_grad = torch::zeros({B, H1, W1, C}, opts);
+    auto fmap2_grad = torch::zeros({B, H2, W2, C}, opts);
+    auto coords_grad = torch::zeros({B, N, H1, W1, 2}, opts);
+
+    const dim3 blocks(B, (H1 + BLOCK_H - 1) / BLOCK_H, (W1 + BLOCK_W - 1) / BLOCK_W);
+    const dim3 threads(BLOCK_H, BLOCK_W);
+
+    altcorr_backward_kernel<float>
+        <<<blocks, threads>>>(fmap1.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                              fmap2.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                              coords.packed_accessor32<float, 5, torch::RestrictPtrTraits>(),
+                              corr_grad.packed_accessor32<float, 5, torch::RestrictPtrTraits>(),
+                              fmap1_grad.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                              fmap2_grad.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                              coords_grad.packed_accessor32<float, 5, torch::RestrictPtrTraits>(), radius);
+
+    return {fmap1_grad, fmap2_grad, coords_grad};
+}

third_party/vipe/csrc/droid_net_ext/correlation_kernels.cu

@@ -0,0 +1,159 @@
+/**
+ * This file includes code originally from the DROID-SLAM repository:
+ * https://github.com/princeton-vl/DROID-SLAM
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <ATen/ATen.h>
+#include <ATen/NativeFunctions.h>
+#include <ATen/Parallel.h>
+
+#define BLOCK 16
+
+__forceinline__ __device__ bool within_bounds(int h, int w, int H, int W) { return h >= 0 && h < H && w >= 0 && w < W; }
+
+template <typename scalar_t>
+__global__ void corr_index_forward_kernel(
+    const torch::PackedTensorAccessor64<scalar_t, 5, torch::RestrictPtrTraits> volume,
+    const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> coords,
+    torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr, int r) {
+    // batch index
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    const int n = blockIdx.z;
+
+    const int h1 = volume.size(1);
+    const int w1 = volume.size(2);
+    const int h2 = volume.size(3);
+    const int w2 = volume.size(4);
+
+    if (!within_bounds(y, x, h1, w1)) {
+        return;
+    }
+
+    float x0 = coords[n][0][y][x];
+    float y0 = coords[n][1][y][x];
+
+    float dx = x0 - floor(x0);
+    float dy = y0 - floor(y0);
+
+    int rd = 2 * r + 1;
+    for (int i = 0; i < rd + 1; i++) {
+        for (int j = 0; j < rd + 1; j++) {
+            int x1 = static_cast<int>(floor(x0)) - r + i;
+            int y1 = static_cast<int>(floor(y0)) - r + j;
+
+            if (within_bounds(y1, x1, h2, w2)) {
+                scalar_t s = volume[n][y][x][y1][x1];
+
+                if (i > 0 && j > 0) corr[n][i - 1][j - 1][y][x] += s * scalar_t(dx * dy);
+
+                if (i > 0 && j < rd) corr[n][i - 1][j][y][x] += s * scalar_t(dx * (1.0f - dy));
+
+                if (i < rd && j > 0) corr[n][i][j - 1][y][x] += s * scalar_t((1.0f - dx) * dy);
+
+                if (i < rd && j < rd) corr[n][i][j][y][x] += s * scalar_t((1.0f - dx) * (1.0f - dy));
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void corr_index_backward_kernel(
+    const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> coords,
+    const torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr_grad,
+    torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> volume_grad, int r) {
+    // batch index
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    const int n = blockIdx.z;
+
+    const int h1 = volume_grad.size(1);
+    const int w1 = volume_grad.size(2);
+    const int h2 = volume_grad.size(3);
+    const int w2 = volume_grad.size(4);
+
+    if (!within_bounds(y, x, h1, w1)) {
+        return;
+    }
+
+    float x0 = coords[n][0][y][x];
+    float y0 = coords[n][1][y][x];
+
+    float dx = x0 - floor(x0);
+    float dy = y0 - floor(y0);
+
+    int rd = 2 * r + 1;
+    for (int i = 0; i < rd + 1; i++) {
+        for (int j = 0; j < rd + 1; j++) {
+            int x1 = static_cast<int>(floor(x0)) - r + i;
+            int y1 = static_cast<int>(floor(y0)) - r + j;
+
+            if (within_bounds(y1, x1, h2, w2)) {
+                scalar_t g = 0.0;
+                if (i > 0 && j > 0) g += corr_grad[n][i - 1][j - 1][y][x] * scalar_t(dx * dy);
+
+                if (i > 0 && j < rd) g += corr_grad[n][i - 1][j][y][x] * scalar_t(dx * (1.0f - dy));
+
+                if (i < rd && j > 0) g += corr_grad[n][i][j - 1][y][x] * scalar_t((1.0f - dx) * dy);
+
+                if (i < rd && j < rd) g += corr_grad[n][i][j][y][x] * scalar_t((1.0f - dx) * (1.0f - dy));
+
+                volume_grad[n][y][x][y1][x1] += g;
+            }
+        }
+    }
+}
+
+std::vector<torch::Tensor> corr_index_cuda_forward(torch::Tensor volume, torch::Tensor coords, int radius) {
+    const auto batch_size = volume.size(0);
+    const auto ht = volume.size(1);
+    const auto wd = volume.size(2);
+
+    const dim3 blocks((wd + BLOCK - 1) / BLOCK, (ht + BLOCK - 1) / BLOCK, batch_size);
+
+    const dim3 threads(BLOCK, BLOCK);
+
+    auto opts = volume.options();
+    torch::Tensor corr = torch::zeros({batch_size, 2 * radius + 1, 2 * radius + 1, ht, wd}, opts);
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(volume.scalar_type(), "sampler_forward_kernel", ([&] {
+                                            corr_index_forward_kernel<scalar_t><<<blocks, threads>>>(
+                                                volume.packed_accessor64<scalar_t, 5, torch::RestrictPtrTraits>(),
+                                                coords.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                                                corr.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
+                                                radius);
+                                        }));
+
+    return {corr};
+}
+
+std::vector<torch::Tensor> corr_index_cuda_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
+                                                    int radius) {
+    const auto batch_size = volume.size(0);
+    const auto ht = volume.size(1);
+    const auto wd = volume.size(2);
+
+    auto volume_grad = torch::zeros_like(volume);
+
+    const dim3 blocks((wd + BLOCK - 1) / BLOCK, (ht + BLOCK - 1) / BLOCK, batch_size);
+
+    const dim3 threads(BLOCK, BLOCK);
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(volume.scalar_type(), "sampler_backward_kernel", ([&] {
+                                            corr_index_backward_kernel<scalar_t><<<blocks, threads>>>(
+                                                coords.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                                                corr_grad.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
+                                                volume_grad.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
+                                                radius);
+                                        }));
+
+    return {volume_grad};
+}

third_party/vipe/csrc/droid_net_ext/droid.cpp

@@ -0,0 +1,63 @@
+/**
+ * This file includes code originally from the DROID-SLAM repository:
+ * https://github.com/princeton-vl/DROID-SLAM
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <vector>
+
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> corr_index_cuda_forward(torch::Tensor volume, torch::Tensor coords, int radius);
+std::vector<torch::Tensor> corr_index_cuda_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
+                                                    int radius);
+std::vector<torch::Tensor> altcorr_cuda_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
+                                                int radius);
+std::vector<torch::Tensor> altcorr_cuda_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
+                                                 torch::Tensor corr_grad, int radius);
+
+// c++ python binding
+std::vector<torch::Tensor> corr_index_forward(torch::Tensor volume, torch::Tensor coords, int radius) {
+    CHECK_INPUT(volume);
+    CHECK_INPUT(coords);
+
+    return corr_index_cuda_forward(volume, coords, radius);
+}
+
+std::vector<torch::Tensor> corr_index_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
+                                               int radius) {
+    CHECK_INPUT(volume);
+    CHECK_INPUT(coords);
+    CHECK_INPUT(corr_grad);
+
+    auto volume_grad = corr_index_cuda_backward(volume, coords, corr_grad, radius);
+    return {volume_grad};
+}
+
+std::vector<torch::Tensor> altcorr_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords, int radius) {
+    CHECK_INPUT(fmap1);
+    CHECK_INPUT(fmap2);
+    CHECK_INPUT(coords);
+
+    return altcorr_cuda_forward(fmap1, fmap2, coords, radius);
+}
+
+std::vector<torch::Tensor> altcorr_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
+                                            torch::Tensor corr_grad, int radius) {
+    CHECK_INPUT(fmap1);
+    CHECK_INPUT(fmap2);
+    CHECK_INPUT(coords);
+    CHECK_INPUT(corr_grad);
+
+    return altcorr_cuda_backward(fmap1, fmap2, coords, corr_grad, radius);
+}
+
+void pybind_droid_net_ext(py::module& m) {
+    // correlation volume kernels
+    m.def("altcorr_forward", &altcorr_forward, "ALTCORR forward");
+    m.def("altcorr_backward", &altcorr_backward, "ALTCORR backward");
+    m.def("corr_index_forward", &corr_index_forward, "INDEX forward");
+    m.def("corr_index_backward", &corr_index_backward, "INDEX backward");
+}

third_party/vipe/csrc/grounding_dino_ext/ms_deform_attn_cuda.cu

@@ -0,0 +1,136 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#include <vector>
+#include "ms_deform_im2col_cuda.cuh"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+namespace groundingdino {
+
+at::Tensor ms_deform_attn_cuda_forward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                       const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
+                                       const at::Tensor &attn_weight, const int im2col_step) {
+    AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+    AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous");
+    AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous");
+    AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
+    AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
+
+    AT_ASSERTM(value.is_cuda(), "value must be a CUDA tensor");
+    AT_ASSERTM(spatial_shapes.is_cuda(), "spatial_shapes must be a CUDA tensor");
+    AT_ASSERTM(level_start_index.is_cuda(), "level_start_index must be a CUDA tensor");
+    AT_ASSERTM(sampling_loc.is_cuda(), "sampling_loc must be a CUDA tensor");
+    AT_ASSERTM(attn_weight.is_cuda(), "attn_weight must be a CUDA tensor");
+
+    const int batch = value.size(0);
+    const int spatial_size = value.size(1);
+    const int num_heads = value.size(2);
+    const int channels = value.size(3);
+
+    const int num_levels = spatial_shapes.size(0);
+
+    const int num_query = sampling_loc.size(1);
+    const int num_point = sampling_loc.size(4);
+
+    const int im2col_step_ = std::min(batch, im2col_step);
+
+    AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
+
+    auto output = at::zeros({batch, num_query, num_heads, channels}, value.options());
+
+    const int batch_n = im2col_step_;
+    auto output_n = output.view({batch / im2col_step_, batch_n, num_query, num_heads, channels});
+    auto per_value_size = spatial_size * num_heads * channels;
+    auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;
+    auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;
+    for (int n = 0; n < batch / im2col_step_; ++n) {
+        auto columns = output_n.select(0, n);
+        AT_DISPATCH_FLOATING_TYPES(
+            value.scalar_type(), "ms_deform_attn_forward_cuda", ([&] {
+                ms_deformable_im2col_cuda(
+                    at::cuda::getCurrentCUDAStream(), value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+                    spatial_shapes.data<int64_t>(), level_start_index.data<int64_t>(),
+                    sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+                    attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size, batch_n, spatial_size,
+                    num_heads, channels, num_levels, num_query, num_point, columns.data<scalar_t>());
+            }));
+    }
+
+    output = output.view({batch, num_query, num_heads * channels});
+
+    return output;
+}
+
+std::vector<at::Tensor> ms_deform_attn_cuda_backward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                                     const at::Tensor &level_start_index,
+                                                     const at::Tensor &sampling_loc, const at::Tensor &attn_weight,
+                                                     const at::Tensor &grad_output, const int im2col_step) {
+    AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+    AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous");
+    AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous");
+    AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
+    AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
+    AT_ASSERTM(grad_output.is_contiguous(), "grad_output tensor has to be contiguous");
+
+    AT_ASSERTM(value.is_cuda(), "value must be a CUDA tensor");
+    AT_ASSERTM(spatial_shapes.is_cuda(), "spatial_shapes must be a CUDA tensor");
+    AT_ASSERTM(level_start_index.is_cuda(), "level_start_index must be a CUDA tensor");
+    AT_ASSERTM(sampling_loc.is_cuda(), "sampling_loc must be a CUDA tensor");
+    AT_ASSERTM(attn_weight.is_cuda(), "attn_weight must be a CUDA tensor");
+    AT_ASSERTM(grad_output.is_cuda(), "grad_output must be a CUDA tensor");
+
+    const int batch = value.size(0);
+    const int spatial_size = value.size(1);
+    const int num_heads = value.size(2);
+    const int channels = value.size(3);
+
+    const int num_levels = spatial_shapes.size(0);
+
+    const int num_query = sampling_loc.size(1);
+    const int num_point = sampling_loc.size(4);
+
+    const int im2col_step_ = std::min(batch, im2col_step);
+
+    AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
+
+    auto grad_value = at::zeros_like(value);
+    auto grad_sampling_loc = at::zeros_like(sampling_loc);
+    auto grad_attn_weight = at::zeros_like(attn_weight);
+
+    const int batch_n = im2col_step_;
+    auto per_value_size = spatial_size * num_heads * channels;
+    auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;
+    auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;
+    auto grad_output_n = grad_output.view({batch / im2col_step_, batch_n, num_query, num_heads, channels});
+
+    for (int n = 0; n < batch / im2col_step_; ++n) {
+        auto grad_output_g = grad_output_n.select(0, n);
+        AT_DISPATCH_FLOATING_TYPES(value.scalar_type(), "ms_deform_attn_backward_cuda", ([&] {
+                                       ms_deformable_col2im_cuda(
+                                           at::cuda::getCurrentCUDAStream(), grad_output_g.data<scalar_t>(),
+                                           value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+                                           spatial_shapes.data<int64_t>(), level_start_index.data<int64_t>(),
+                                           sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+                                           attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size,
+                                           batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point,
+                                           grad_value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+                                           grad_sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+                                           grad_attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size);
+                                   }));
+    }
+
+    return {grad_value, grad_sampling_loc, grad_attn_weight};
+}
+
+}  // namespace groundingdino

third_party/vipe/csrc/grounding_dino_ext/ms_deform_attn_cuda.h

@@ -0,0 +1,25 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+namespace groundingdino {
+
+at::Tensor ms_deform_attn_cuda_forward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                       const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
+                                       const at::Tensor &attn_weight, const int im2col_step);
+
+std::vector<at::Tensor> ms_deform_attn_cuda_backward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                                     const at::Tensor &level_start_index,
+                                                     const at::Tensor &sampling_loc, const at::Tensor &attn_weight,
+                                                     const at::Tensor &grad_output, const int im2col_step);
+
+}  // namespace groundingdino

third_party/vipe/csrc/grounding_dino_ext/ms_deform_im2col_cuda.cuh

@@ -0,0 +1,914 @@
+/*!
+**************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************
+* Modified from DCN (https://github.com/msracver/Deformable-ConvNets)
+* Copyright (c) 2018 Microsoft
+**************************************************************************
+*/
+
+#include <cstdio>
+#include <algorithm>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#define CUDA_KERNEL_LOOP(i, n) for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x)
+
+const int CUDA_NUM_THREADS = 1024;
+inline int GET_BLOCKS(const int N, const int num_threads) { return (N + num_threads - 1) / num_threads; }
+
+template <typename scalar_t>
+__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t *&bottom_data, const int &height, const int &width,
+                                                   const int &nheads, const int &channels, const scalar_t &h,
+                                                   const scalar_t &w, const int &m, const int &c) {
+    const int h_low = floor(h);
+    const int w_low = floor(w);
+    const int h_high = h_low + 1;
+    const int w_high = w_low + 1;
+
+    const scalar_t lh = h - h_low;
+    const scalar_t lw = w - w_low;
+    const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+    const int w_stride = nheads * channels;
+    const int h_stride = width * w_stride;
+    const int h_low_ptr_offset = h_low * h_stride;
+    const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+    const int w_low_ptr_offset = w_low * w_stride;
+    const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+    const int base_ptr = m * channels + c;
+
+    scalar_t v1 = 0;
+    if (h_low >= 0 && w_low >= 0) {
+        const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+        v1 = bottom_data[ptr1];
+    }
+    scalar_t v2 = 0;
+    if (h_low >= 0 && w_high <= width - 1) {
+        const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+        v2 = bottom_data[ptr2];
+    }
+    scalar_t v3 = 0;
+    if (h_high <= height - 1 && w_low >= 0) {
+        const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+        v3 = bottom_data[ptr3];
+    }
+    scalar_t v4 = 0;
+    if (h_high <= height - 1 && w_high <= width - 1) {
+        const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+        v4 = bottom_data[ptr4];
+    }
+
+    const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+
+    const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+    return val;
+}
+
+template <typename scalar_t>
+__device__ void ms_deform_attn_col2im_bilinear(const scalar_t *&bottom_data, const int &height, const int &width,
+                                               const int &nheads, const int &channels, const scalar_t &h,
+                                               const scalar_t &w, const int &m, const int &c, const scalar_t &top_grad,
+                                               const scalar_t &attn_weight, scalar_t *&grad_value,
+                                               scalar_t *grad_sampling_loc, scalar_t *grad_attn_weight) {
+    const int h_low = floor(h);
+    const int w_low = floor(w);
+    const int h_high = h_low + 1;
+    const int w_high = w_low + 1;
+
+    const scalar_t lh = h - h_low;
+    const scalar_t lw = w - w_low;
+    const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+    const int w_stride = nheads * channels;
+    const int h_stride = width * w_stride;
+    const int h_low_ptr_offset = h_low * h_stride;
+    const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+    const int w_low_ptr_offset = w_low * w_stride;
+    const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+    const int base_ptr = m * channels + c;
+
+    const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+    const scalar_t top_grad_value = top_grad * attn_weight;
+    scalar_t grad_h_weight = 0, grad_w_weight = 0;
+
+    scalar_t v1 = 0;
+    if (h_low >= 0 && w_low >= 0) {
+        const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+        v1 = bottom_data[ptr1];
+        grad_h_weight -= hw * v1;
+        grad_w_weight -= hh * v1;
+        atomicAdd(grad_value + ptr1, w1 * top_grad_value);
+    }
+    scalar_t v2 = 0;
+    if (h_low >= 0 && w_high <= width - 1) {
+        const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+        v2 = bottom_data[ptr2];
+        grad_h_weight -= lw * v2;
+        grad_w_weight += hh * v2;
+        atomicAdd(grad_value + ptr2, w2 * top_grad_value);
+    }
+    scalar_t v3 = 0;
+    if (h_high <= height - 1 && w_low >= 0) {
+        const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+        v3 = bottom_data[ptr3];
+        grad_h_weight += hw * v3;
+        grad_w_weight -= lh * v3;
+        atomicAdd(grad_value + ptr3, w3 * top_grad_value);
+    }
+    scalar_t v4 = 0;
+    if (h_high <= height - 1 && w_high <= width - 1) {
+        const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+        v4 = bottom_data[ptr4];
+        grad_h_weight += lw * v4;
+        grad_w_weight += lh * v4;
+        atomicAdd(grad_value + ptr4, w4 * top_grad_value);
+    }
+
+    const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+    *grad_attn_weight = top_grad * val;
+    *grad_sampling_loc = width * grad_w_weight * top_grad_value;
+    *(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value;
+}
+
+template <typename scalar_t>
+__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t *&bottom_data, const int &height, const int &width,
+                                                  const int &nheads, const int &channels, const scalar_t &h,
+                                                  const scalar_t &w, const int &m, const int &c,
+                                                  const scalar_t &top_grad, const scalar_t &attn_weight,
+                                                  scalar_t *&grad_value, scalar_t *grad_sampling_loc,
+                                                  scalar_t *grad_attn_weight) {
+    const int h_low = floor(h);
+    const int w_low = floor(w);
+    const int h_high = h_low + 1;
+    const int w_high = w_low + 1;
+
+    const scalar_t lh = h - h_low;
+    const scalar_t lw = w - w_low;
+    const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+    const int w_stride = nheads * channels;
+    const int h_stride = width * w_stride;
+    const int h_low_ptr_offset = h_low * h_stride;
+    const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+    const int w_low_ptr_offset = w_low * w_stride;
+    const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+    const int base_ptr = m * channels + c;
+
+    const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+    const scalar_t top_grad_value = top_grad * attn_weight;
+    scalar_t grad_h_weight = 0, grad_w_weight = 0;
+
+    scalar_t v1 = 0;
+    if (h_low >= 0 && w_low >= 0) {
+        const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+        v1 = bottom_data[ptr1];
+        grad_h_weight -= hw * v1;
+        grad_w_weight -= hh * v1;
+        atomicAdd(grad_value + ptr1, w1 * top_grad_value);
+    }
+    scalar_t v2 = 0;
+    if (h_low >= 0 && w_high <= width - 1) {
+        const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+        v2 = bottom_data[ptr2];
+        grad_h_weight -= lw * v2;
+        grad_w_weight += hh * v2;
+        atomicAdd(grad_value + ptr2, w2 * top_grad_value);
+    }
+    scalar_t v3 = 0;
+    if (h_high <= height - 1 && w_low >= 0) {
+        const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+        v3 = bottom_data[ptr3];
+        grad_h_weight += hw * v3;
+        grad_w_weight -= lh * v3;
+        atomicAdd(grad_value + ptr3, w3 * top_grad_value);
+    }
+    scalar_t v4 = 0;
+    if (h_high <= height - 1 && w_high <= width - 1) {
+        const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+        v4 = bottom_data[ptr4];
+        grad_h_weight += lw * v4;
+        grad_w_weight += lh * v4;
+        atomicAdd(grad_value + ptr4, w4 * top_grad_value);
+    }
+
+    const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+    atomicAdd(grad_attn_weight, top_grad * val);
+    atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value);
+    atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value);
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_im2col_gpu_kernel(const int n, const scalar_t *data_value,
+                                                const int64_t *data_spatial_shapes,
+                                                const int64_t *data_level_start_index,
+                                                const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+                                                const int batch_size, const int spatial_size, const int num_heads,
+                                                const int channels, const int num_levels, const int num_query,
+                                                const int num_point, scalar_t *data_col) {
+    CUDA_KERNEL_LOOP(index, n) {
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        scalar_t *data_col_ptr = data_col + index;
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+        scalar_t col = 0;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride);
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels,
+                                                          h_im, w_im, m_col, c_col) *
+                           weight;
+                }
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+            }
+        }
+        *data_col_ptr = col;
+    }
+}
+
+template <typename scalar_t, unsigned int blockSize>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(
+    const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+    const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+    const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
+    const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
+    scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];
+        __shared__ scalar_t cache_grad_attn_weight[blockSize];
+        unsigned int tid = threadIdx.x;
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                *(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
+                *(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
+                *(cache_grad_attn_weight + threadIdx.x) = 0;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                   w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                   cache_grad_sampling_loc + (threadIdx.x << 1),
+                                                   cache_grad_attn_weight + threadIdx.x);
+                }
+
+                __syncthreads();
+                if (tid == 0) {
+                    scalar_t _grad_w = cache_grad_sampling_loc[0], _grad_h = cache_grad_sampling_loc[1],
+                             _grad_a = cache_grad_attn_weight[0];
+                    int sid = 2;
+                    for (unsigned int tid = 1; tid < blockSize; ++tid) {
+                        _grad_w += cache_grad_sampling_loc[sid];
+                        _grad_h += cache_grad_sampling_loc[sid + 1];
+                        _grad_a += cache_grad_attn_weight[tid];
+                        sid += 2;
+                    }
+
+                    *grad_sampling_loc = _grad_w;
+                    *(grad_sampling_loc + 1) = _grad_h;
+                    *grad_attn_weight = _grad_a;
+                }
+                __syncthreads();
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t, unsigned int blockSize>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(
+    const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+    const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+    const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
+    const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
+    scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];
+        __shared__ scalar_t cache_grad_attn_weight[blockSize];
+        unsigned int tid = threadIdx.x;
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                *(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
+                *(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
+                *(cache_grad_attn_weight + threadIdx.x) = 0;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                   w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                   cache_grad_sampling_loc + (threadIdx.x << 1),
+                                                   cache_grad_attn_weight + threadIdx.x);
+                }
+
+                __syncthreads();
+
+                for (unsigned int s = blockSize / 2; s > 0; s >>= 1) {
+                    if (tid < s) {
+                        const unsigned int xid1 = tid << 1;
+                        const unsigned int xid2 = (tid + s) << 1;
+                        cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+                        cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+                        cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+                    }
+                    __syncthreads();
+                }
+
+                if (tid == 0) {
+                    *grad_sampling_loc = cache_grad_sampling_loc[0];
+                    *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];
+                    *grad_attn_weight = cache_grad_attn_weight[0];
+                }
+                __syncthreads();
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(
+    const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+    const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+    const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
+    const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
+    scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        extern __shared__ int _s[];
+        scalar_t *cache_grad_sampling_loc = (scalar_t *)_s;
+        scalar_t *cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+        unsigned int tid = threadIdx.x;
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                *(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
+                *(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
+                *(cache_grad_attn_weight + threadIdx.x) = 0;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                   w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                   cache_grad_sampling_loc + (threadIdx.x << 1),
+                                                   cache_grad_attn_weight + threadIdx.x);
+                }
+
+                __syncthreads();
+                if (tid == 0) {
+                    scalar_t _grad_w = cache_grad_sampling_loc[0], _grad_h = cache_grad_sampling_loc[1],
+                             _grad_a = cache_grad_attn_weight[0];
+                    int sid = 2;
+                    for (unsigned int tid = 1; tid < blockDim.x; ++tid) {
+                        _grad_w += cache_grad_sampling_loc[sid];
+                        _grad_h += cache_grad_sampling_loc[sid + 1];
+                        _grad_a += cache_grad_attn_weight[tid];
+                        sid += 2;
+                    }
+
+                    *grad_sampling_loc = _grad_w;
+                    *(grad_sampling_loc + 1) = _grad_h;
+                    *grad_attn_weight = _grad_a;
+                }
+                __syncthreads();
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(
+    const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+    const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+    const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
+    const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
+    scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        extern __shared__ int _s[];
+        scalar_t *cache_grad_sampling_loc = (scalar_t *)_s;
+        scalar_t *cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+        unsigned int tid = threadIdx.x;
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                *(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
+                *(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
+                *(cache_grad_attn_weight + threadIdx.x) = 0;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                   w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                   cache_grad_sampling_loc + (threadIdx.x << 1),
+                                                   cache_grad_attn_weight + threadIdx.x);
+                }
+
+                __syncthreads();
+
+                for (unsigned int s = blockDim.x / 2, spre = blockDim.x; s > 0; s >>= 1, spre >>= 1) {
+                    if (tid < s) {
+                        const unsigned int xid1 = tid << 1;
+                        const unsigned int xid2 = (tid + s) << 1;
+                        cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+                        cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+                        cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+                        if (tid + (s << 1) < spre) {
+                            cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];
+                            cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];
+                            cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];
+                        }
+                    }
+                    __syncthreads();
+                }
+
+                if (tid == 0) {
+                    *grad_sampling_loc = cache_grad_sampling_loc[0];
+                    *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];
+                    *grad_attn_weight = cache_grad_attn_weight[0];
+                }
+                __syncthreads();
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(
+    const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+    const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+    const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
+    const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
+    scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        extern __shared__ int _s[];
+        scalar_t *cache_grad_sampling_loc = (scalar_t *)_s;
+        scalar_t *cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+        unsigned int tid = threadIdx.x;
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                *(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
+                *(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
+                *(cache_grad_attn_weight + threadIdx.x) = 0;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                   w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                   cache_grad_sampling_loc + (threadIdx.x << 1),
+                                                   cache_grad_attn_weight + threadIdx.x);
+                }
+
+                __syncthreads();
+
+                for (unsigned int s = blockDim.x / 2, spre = blockDim.x; s > 0; s >>= 1, spre >>= 1) {
+                    if (tid < s) {
+                        const unsigned int xid1 = tid << 1;
+                        const unsigned int xid2 = (tid + s) << 1;
+                        cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+                        cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+                        cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+                        if (tid + (s << 1) < spre) {
+                            cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];
+                            cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];
+                            cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];
+                        }
+                    }
+                    __syncthreads();
+                }
+
+                if (tid == 0) {
+                    atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]);
+                    atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]);
+                    atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]);
+                }
+                __syncthreads();
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n, const scalar_t *grad_col, const scalar_t *data_value,
+                                                   const int64_t *data_spatial_shapes,
+                                                   const int64_t *data_level_start_index,
+                                                   const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+                                                   const int batch_size, const int spatial_size, const int num_heads,
+                                                   const int channels, const int num_levels, const int num_query,
+                                                   const int num_point, scalar_t *grad_value,
+                                                   scalar_t *grad_sampling_loc, scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear_gm(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                      w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                      grad_sampling_loc, grad_attn_weight);
+                }
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+void ms_deformable_im2col_cuda(cudaStream_t stream, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+                               const int64_t *data_level_start_index, const scalar_t *data_sampling_loc,
+                               const scalar_t *data_attn_weight, const int batch_size, const int spatial_size,
+                               const int num_heads, const int channels, const int num_levels, const int num_query,
+                               const int num_point, scalar_t *data_col) {
+    const int num_kernels = batch_size * num_query * num_heads * channels;
+    const int num_actual_kernels = batch_size * num_query * num_heads * channels;
+    const int num_threads = CUDA_NUM_THREADS;
+    ms_deformable_im2col_gpu_kernel<scalar_t><<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+        num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight,
+        batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col);
+
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess) {
+        printf("error in ms_deformable_im2col_cuda: %s\n", cudaGetErrorString(err));
+    }
+}
+
+template <typename scalar_t>
+void ms_deformable_col2im_cuda(cudaStream_t stream, const scalar_t *grad_col, const scalar_t *data_value,
+                               const int64_t *data_spatial_shapes, const int64_t *data_level_start_index,
+                               const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+                               const int batch_size, const int spatial_size, const int num_heads, const int channels,
+                               const int num_levels, const int num_query, const int num_point, scalar_t *grad_value,
+                               scalar_t *grad_sampling_loc, scalar_t *grad_attn_weight) {
+    const int num_threads = (channels > CUDA_NUM_THREADS) ? CUDA_NUM_THREADS : channels;
+    const int num_kernels = batch_size * num_query * num_heads * channels;
+    const int num_actual_kernels = batch_size * num_query * num_heads * channels;
+    if (channels > 1024) {
+        if ((channels & 1023) == 0) {
+            ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks<scalar_t>
+                <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, num_threads * 3 * sizeof(scalar_t),
+                   stream>>>(num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                             data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels,
+                             num_levels, num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+        } else {
+            ms_deformable_col2im_gpu_kernel_gm<scalar_t>
+                <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                    num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc,
+                    data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point,
+                    grad_value, grad_sampling_loc, grad_attn_weight);
+        }
+    } else {
+        switch (channels) {
+            case 1:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 1>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 2:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 2>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 4:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 4>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 8:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 8>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 16:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 16>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 32:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 32>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 64:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 64>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 128:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 128>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 256:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 256>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 512:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 512>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 1024:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 1024>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            default:
+                if (channels < 64) {
+                    ms_deformable_col2im_gpu_kernel_shm_reduce_v1<scalar_t>
+                        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, num_threads * 3 * sizeof(scalar_t),
+                           stream>>>(num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                                     data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels,
+                                     num_levels, num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                } else {
+                    ms_deformable_col2im_gpu_kernel_shm_reduce_v2<scalar_t>
+                        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, num_threads * 3 * sizeof(scalar_t),
+                           stream>>>(num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                                     data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels,
+                                     num_levels, num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                }
+        }
+    }
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess) {
+        printf("error in ms_deformable_col2im_cuda: %s\n", cudaGetErrorString(err));
+    }
+}

third_party/vipe/csrc/grounding_dino_ext/vision.cpp

@@ -0,0 +1,33 @@
+// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+#include "ms_deform_attn_cuda.h"
+
+namespace groundingdino {
+
+at::Tensor ms_deform_attn_forward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                  const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
+                                  const at::Tensor &attn_weight, const int im2col_step) {
+    if (value.is_cuda()) {
+        return ms_deform_attn_cuda_forward(value, spatial_shapes, level_start_index, sampling_loc, attn_weight,
+                                           im2col_step);
+    }
+    AT_ERROR("Not implemented on the CPU");
+}
+
+std::vector<at::Tensor> ms_deform_attn_backward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                                const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
+                                                const at::Tensor &attn_weight, const at::Tensor &grad_output,
+                                                const int im2col_step) {
+    if (value.is_cuda()) {
+        return ms_deform_attn_cuda_backward(value, spatial_shapes, level_start_index, sampling_loc, attn_weight,
+                                            grad_output, im2col_step);
+    }
+    AT_ERROR("Not implemented on the CPU");
+}
+
+}  // namespace groundingdino
+
+void pybind_grounding_dino_ext(py::module &m) {
+    m.def("ms_deform_attn_forward", &groundingdino::ms_deform_attn_forward, "ms_deform_attn_forward");
+    m.def("ms_deform_attn_backward", &groundingdino::ms_deform_attn_backward, "ms_deform_attn_backward");
+}

third_party/vipe/csrc/lietorch_ext/common.h

@@ -0,0 +1,16 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef COMMON_H
+#define COMMON_H
+
+#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int
+#define EIGEN_RUNTIME_NO_MALLOC
+
+#define EPS 1e-6
+#define PI 3.14159265358979323846
+
+#endif

third_party/vipe/csrc/lietorch_ext/dispatch.h

@@ -0,0 +1,52 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef DISPATCH_H
+#define DISPATCH_H
+
+#include <torch/extension.h>
+
+#include "so3.h"
+#include "rxso3.h"
+#include "se3.h"
+#include "sim3.h"
+
+#define PRIVATE_CASE_TYPE(group_index, enum_type, type, ...) \
+    case enum_type: {                                        \
+        using scalar_t = type;                               \
+        switch (group_index) {                               \
+            case 1: {                                        \
+                using group_t = SO3<type>;                   \
+                return __VA_ARGS__();                        \
+            }                                                \
+            case 2: {                                        \
+                using group_t = RxSO3<type>;                 \
+                return __VA_ARGS__();                        \
+            }                                                \
+            case 3: {                                        \
+                using group_t = SE3<type>;                   \
+                return __VA_ARGS__();                        \
+            }                                                \
+            case 4: {                                        \
+                using group_t = Sim3<type>;                  \
+                return __VA_ARGS__();                        \
+            }                                                \
+        }                                                    \
+    }
+
+#define DISPATCH_GROUP_AND_FLOATING_TYPES(GROUP_INDEX, TYPE, NAME, ...)                 \
+    [&] {                                                                               \
+        const auto& the_type = TYPE;                                                    \
+        /* don't use TYPE again in case it is an expensive or side-effect op */         \
+        switch (the_type) {                                                             \
+            PRIVATE_CASE_TYPE(GROUP_INDEX, at::ScalarType::Double, double, __VA_ARGS__) \
+            PRIVATE_CASE_TYPE(GROUP_INDEX, at::ScalarType::Float, float, __VA_ARGS__)   \
+            default:                                                                    \
+                break;                                                                  \
+        }                                                                               \
+    }()
+
+#endif

third_party/vipe/csrc/lietorch_ext/lietorch.cpp

@@ -0,0 +1,336 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <vector>
+#include "lietorch_gpu.h"
+#include "lietorch_cpu.h"
+#include <c10/cuda/CUDAGuard.h>
+
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+
+/* Interface for cuda and c++ group operations
+
+    enum group_t { SO3=1, SE3=2, Sim3=3 };
+    X, Y, Z: (uppercase) Lie Group Elements
+    a, b, c: (lowercase) Lie Algebra Elements
+*/
+
+// Unary operations
+torch::Tensor expm(int group_index, torch::Tensor a) {
+    CHECK_CONTIGUOUS(a);
+    if (a.device().type() == torch::DeviceType::CPU) {
+        return exp_forward_cpu(group_index, a);
+
+    } else if (a.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+        return exp_forward_gpu(group_index, a);
+    }
+
+    return a;
+}
+
+std::vector<torch::Tensor> expm_backward(int group_index, torch::Tensor grad, torch::Tensor a) {
+    CHECK_CONTIGUOUS(a);
+    CHECK_CONTIGUOUS(grad);
+    if (a.device().type() == torch::DeviceType::CPU) {
+        return exp_backward_cpu(group_index, grad, a);
+
+    } else if (a.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+        return exp_backward_gpu(group_index, grad, a);
+    }
+
+    return {};
+}
+
+torch::Tensor logm(int group_index, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return log_forward_cpu(group_index, X);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return log_forward_gpu(group_index, X);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> logm_backward(int group_index, torch::Tensor grad, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return log_backward_cpu(group_index, grad, X);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return log_backward_gpu(group_index, grad, X);
+    }
+
+    return {};
+}
+
+torch::Tensor inv(int group_index, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return inv_forward_cpu(group_index, X);
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return inv_forward_gpu(group_index, X);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> inv_backward(int group_index, torch::Tensor grad, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return inv_backward_cpu(group_index, grad, X);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return inv_backward_gpu(group_index, grad, X);
+    }
+
+    return {};
+}
+
+// Binary operations
+
+torch::Tensor mul(int group_index, torch::Tensor X, torch::Tensor Y) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(Y);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return mul_forward_cpu(group_index, X, Y);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return mul_forward_gpu(group_index, X, Y);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> mul_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(Y);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return mul_backward_cpu(group_index, grad, X, Y);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return mul_backward_gpu(group_index, grad, X, Y);
+    }
+
+    return {};
+}
+
+torch::Tensor adj(int group_index, torch::Tensor X, torch::Tensor a) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(a);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return adj_forward_cpu(group_index, X, a);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return adj_forward_gpu(group_index, X, a);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> adj_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(a);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return adj_backward_cpu(group_index, grad, X, a);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return adj_backward_gpu(group_index, grad, X, a);
+    }
+
+    return {};
+}
+
+torch::Tensor adjT(int group_index, torch::Tensor X, torch::Tensor a) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(a);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return adjT_forward_cpu(group_index, X, a);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return adjT_forward_gpu(group_index, X, a);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> adjT_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(a);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return adjT_backward_cpu(group_index, grad, X, a);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return adjT_backward_gpu(group_index, grad, X, a);
+    }
+
+    return {};
+}
+
+torch::Tensor act(int group_index, torch::Tensor X, torch::Tensor p) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(p);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return act_forward_cpu(group_index, X, p);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return act_forward_gpu(group_index, X, p);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> act_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(p);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return act_backward_cpu(group_index, grad, X, p);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return act_backward_gpu(group_index, grad, X, p);
+    }
+
+    return {};
+}
+
+torch::Tensor act4(int group_index, torch::Tensor X, torch::Tensor p) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(p);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return act4_forward_cpu(group_index, X, p);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return act4_forward_gpu(group_index, X, p);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> act4_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(p);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return act4_backward_cpu(group_index, grad, X, p);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return act4_backward_gpu(group_index, grad, X, p);
+    }
+
+    return {};
+}
+
+torch::Tensor projector(int group_index, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return orthogonal_projector_cpu(group_index, X);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return orthogonal_projector_gpu(group_index, X);
+    }
+
+    return X;
+}
+
+torch::Tensor as_matrix(int group_index, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return as_matrix_forward_cpu(group_index, X);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return as_matrix_forward_gpu(group_index, X);
+    }
+
+    return X;
+}
+
+torch::Tensor Jinv(int group_index, torch::Tensor X, torch::Tensor a) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(a);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return jleft_forward_cpu(group_index, X, a);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return jleft_forward_gpu(group_index, X, a);
+    }
+
+    return a;
+}
+
+void pybind_lietorch_ext(py::module &m) {
+    // {exp, log, inv, mul, adj, adjT, act, act4} forward/backward bindings
+    m.def("expm", &expm, "exp map forward");
+    m.def("expm_backward", &expm_backward, "exp map backward");
+
+    m.def("logm", &logm, "log map forward");
+    m.def("logm_backward", &logm_backward, "log map backward");
+
+    m.def("inv", &inv, "inverse operator");
+    m.def("inv_backward", &inv_backward, "inverse operator backward");
+
+    m.def("mul", &mul, "group operator");
+    m.def("mul_backward", &mul_backward, "group operator backward");
+
+    m.def("adj", &adj, "adjoint operator");
+    m.def("adj_backward", &adj_backward, "adjoint operator backward");
+
+    m.def("adjT", &adjT, "transposed adjoint operator");
+    m.def("adjT_backward", &adjT_backward, "transposed adjoint operator backward");
+
+    m.def("act", &act, "action on point");
+    m.def("act_backward", &act_backward, "action on point backward");
+
+    m.def("act4", &act4, "action on homogeneous point");
+    m.def("act4_backward", &act4_backward, "action on homogeneous point backward");
+
+    // functions with no gradient
+    m.def("as_matrix", &as_matrix, "convert to matrix");
+    m.def("projector", &projector, "orthogonal projection matrix");
+    m.def("Jinv", &Jinv, "left inverse jacobian operator");
+}

third_party/vipe/csrc/lietorch_ext/lietorch_cpu.cpp

@@ -0,0 +1,610 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include "lietorch_cpu.h"
+#include <eigen3/Eigen/Dense>
+
+#include <iostream>
+#include "common.h"
+#include "dispatch.h"
+
+#include "so3.h"
+#include "rxso3.h"
+#include "se3.h"
+#include "sim3.h"
+
+template <typename Group, typename scalar_t>
+void exp_forward_kernel(const scalar_t* a_ptr, scalar_t* X_ptr, int batch_size) {
+    // exponential map forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Tangent a(a_ptr + i * Group::K);
+            Eigen::Map<Data>(X_ptr + i * Group::N) = Group::Exp(a).data();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void exp_backward_kernel(const scalar_t* grad, const scalar_t* a_ptr, scalar_t* da, int batch_size) {
+    // exponential map backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Tangent a(a_ptr + i * Group::K);
+            Grad dX(grad + i * Group::N);
+            Eigen::Map<Grad>(da + i * Group::K) = dX * Group::left_jacobian(a);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void log_forward_kernel(const scalar_t* X_ptr, scalar_t* a_ptr, int batch_size) {
+    // logarithm map forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Tangent a = Group(X_ptr + i * Group::N).Log();
+            Eigen::Map<Tangent>(a_ptr + i * Group::K) = a;
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void log_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
+    // logarithm map backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Tangent a = Group(X_ptr + i * Group::N).Log();
+            Grad da(grad + i * Group::K);
+            Eigen::Map<Grad>(dX + i * Group::N) = da * Group::left_jacobian_inverse(a);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void inv_forward_kernel(const scalar_t* X_ptr, scalar_t* Y_ptr, int batch_size) {
+    // group inverse forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Eigen::Map<Data>(Y_ptr + i * Group::N) = X.inv().data();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void inv_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
+    // group inverse backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group Y = Group(X_ptr + i * Group::N).inv();
+            Grad dY(grad + i * Group::N);
+            Eigen::Map<Grad>(dX + i * Group::N) = -dY * Y.Adj();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void mul_forward_kernel(const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* Z_ptr, int batch_size) {
+    // group multiplication forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group Z = Group(X_ptr + i * Group::N) * Group(Y_ptr + i * Group::N);
+            Eigen::Map<Data>(Z_ptr + i * Group::N) = Z.data();
+        }
+    });
+}
+
+template <class Group, typename scalar_t>
+void mul_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* dX, scalar_t* dY,
+                         int batch_size) {
+    // group multiplication backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Grad dZ(grad + i * Group::N);
+            Group X(X_ptr + i * Group::N);
+            Eigen::Map<Grad>(dX + i * Group::N) = dZ;
+            Eigen::Map<Grad>(dY + i * Group::N) = dZ * X.Adj();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void adj_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
+    // adjoint forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Tangent a(a_ptr + i * Group::K);
+            Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.Adj(a);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void adj_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX, scalar_t* da,
+                         int batch_size) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Grad db(grad + i * Group::K);
+
+            Tangent a(a_ptr + i * Group::K);
+            Tangent b = X.Adj() * a;
+
+            Eigen::Map<Grad>(da + i * Group::K) = db * X.Adj();
+            Eigen::Map<Grad>(dX + i * Group::N) = -db * Group::adj(b);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void adjT_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
+    // adjoint forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Tangent a(a_ptr + i * Group::K);
+            Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.AdjT(a);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void adjT_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX,
+                          scalar_t* da, int batch_size) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Tangent db(grad + i * Group::K);
+            Grad a(a_ptr + i * Group::K);
+
+            Eigen::Map<Tangent>(da + i * Group::K) = X.Adj(db);
+            Eigen::Map<Grad>(dX + i * Group::N) = -a * Group::adj(X.Adj(db));
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void act_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int batch_size) {
+    // action on point forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Point = Eigen::Matrix<scalar_t, 3, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Point p(p_ptr + i * 3);
+            Eigen::Map<Point>(q_ptr + i * 3) = X * p;
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void act_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX, scalar_t* dp,
+                         int batch_size) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Point = Eigen::Matrix<scalar_t, 3, 1>;
+    using PointGrad = Eigen::Matrix<scalar_t, 1, 3>;
+    using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Point p(p_ptr + i * 3);
+            PointGrad dq(grad + i * 3);
+
+            Eigen::Map<PointGrad>(dp + i * 3) = dq * X.Matrix().template block<3, 3>(0, 0);
+            Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act_jacobian(X * p);
+        }
+    });
+}
+
+// template <typename Group, typename scalar_t>
+// void tovec_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
+//     // group inverse forward kernel
+//     using Data = Eigen::Matrix<scalar_t,Group::N,1>;
+//     using Grad = Eigen::Matrix<scalar_t,1,Group::N>;
+
+//     at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+//         for (int64_t i=start; i<end; i++) {
+//             Group X(X_ptr + i*Group::N);
+//             Grad g(grad + i*Group::N);
+//             Eigen::Map<Grad>(dX + i*Group::N) = g * X.vec_jacobian();
+//         }
+//     });
+// }
+
+// template <typename Group, typename scalar_t>
+// void fromvec_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
+//     // group inverse forward kernel
+//     using Data = Eigen::Matrix<scalar_t,Group::N,1>;
+//     using Grad = Eigen::Matrix<scalar_t,1,Group::N>;
+
+//     at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+//         for (int64_t i=start; i<end; i++) {
+//             Group X(X_ptr + i*Group::N);
+//             Grad g(grad + i*Group::N);
+//             Eigen::Map<Grad>(dX + i*Group::N) = g * X.vec_jacobian();
+//         }
+//     });
+// }
+
+template <typename Group, typename scalar_t>
+void act4_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int batch_size) {
+    // action on homogeneous point forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Point = Eigen::Matrix<scalar_t, 4, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Point p(p_ptr + i * 4);
+            Eigen::Map<Point>(q_ptr + i * 4) = X.act4(p);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void act4_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX,
+                          scalar_t* dp, int batch_size) {
+    // action on homogeneous point backward kernel
+
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Point = Eigen::Matrix<scalar_t, 4, 1>;
+    using PointGrad = Eigen::Matrix<scalar_t, 1, 4>;
+    using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Point p(p_ptr + i * 4);
+            PointGrad dq(grad + i * 4);
+
+            Eigen::Map<PointGrad>(dp + i * 4) = dq * X.Matrix4x4();
+            const Point q = X.act4(p);
+            Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act4_jacobian(q);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void as_matrix_forward_kernel(const scalar_t* X_ptr, scalar_t* T_ptr, int batch_size) {
+    // group inverse forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Matrix4 = Eigen::Matrix<scalar_t, 4, 4, Eigen::RowMajor>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Eigen::Map<Matrix4>(T_ptr + i * 16) = X.Matrix4x4();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void orthogonal_projector_kernel(const scalar_t* X_ptr, scalar_t* P_ptr, int batch_size) {
+    // group inverse forward kernel
+    using Proj = Eigen::Matrix<scalar_t, Group::N, Group::N, Eigen::RowMajor>;
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Eigen::Map<Proj>(P_ptr + i * Group::N * Group::N) = X.orthogonal_projector();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void jleft_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
+    // left-jacobian inverse action
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Tangent a(a_ptr + i * Group::K);
+            Tangent b = Group::left_jacobian_inverse(X.Log()) * a;
+            Eigen::Map<Tangent>(b_ptr + i * Group::K) = b;
+        }
+    });
+}
+
+// unary operations
+
+torch::Tensor exp_forward_cpu(int group_id, torch::Tensor a) {
+    int batch_size = a.size(0);
+    torch::Tensor X;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_forward_kernel", ([&] {
+                                          X = torch::zeros({batch_size, group_t::N}, a.options());
+                                          exp_forward_kernel<group_t, scalar_t>(a.data_ptr<scalar_t>(),
+                                                                                X.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return X;
+}
+
+std::vector<torch::Tensor> exp_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor a) {
+    int batch_size = a.size(0);
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_backward_kernel", ([&] {
+                                          exp_backward_kernel<group_t, scalar_t>(grad.data_ptr<scalar_t>(),
+                                                                                 a.data_ptr<scalar_t>(),
+                                                                                 da.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {da};
+}
+
+torch::Tensor log_forward_cpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor a;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_forward_kernel", ([&] {
+                                          a = torch::zeros({batch_size, group_t::K}, X.options());
+                                          log_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                a.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return a;
+}
+
+std::vector<torch::Tensor> log_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_backward_kernel", ([&] {
+                                          log_backward_kernel<group_t, scalar_t>(grad.data_ptr<scalar_t>(),
+                                                                                 X.data_ptr<scalar_t>(),
+                                                                                 dX.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX};
+}
+
+torch::Tensor inv_forward_cpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor Y = torch::zeros_like(X);
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "inv_forward_kernel",
+        ([&] { inv_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), batch_size); }));
+
+    return Y;
+}
+
+std::vector<torch::Tensor> inv_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "inv_backward_kernel", ([&] {
+                                          inv_backward_kernel<group_t, scalar_t>(grad.data_ptr<scalar_t>(),
+                                                                                 X.data_ptr<scalar_t>(),
+                                                                                 dX.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX};
+}
+
+// binary operations
+torch::Tensor mul_forward_cpu(int group_id, torch::Tensor X, torch::Tensor Y) {
+    int batch_size = X.size(0);
+    torch::Tensor Z = torch::zeros_like(X);
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "mul_forward_kernel", ([&] {
+                                          mul_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                Y.data_ptr<scalar_t>(),
+                                                                                Z.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return Z;
+}
+
+std::vector<torch::Tensor> mul_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dY = torch::zeros(Y.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "mul_backward_kernel", ([&] {
+                                          mul_backward_kernel<group_t, scalar_t>(
+                                              grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(),
+                                              dX.data_ptr<scalar_t>(), dY.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX, dY};
+}
+
+torch::Tensor adj_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adj_forward_kernel", ([&] {
+                                          adj_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                a.data_ptr<scalar_t>(),
+                                                                                b.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return b;
+}
+
+std::vector<torch::Tensor> adj_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adj_backward_kernel", ([&] {
+                                          adj_backward_kernel<group_t, scalar_t>(
+                                              grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(),
+                                              dX.data_ptr<scalar_t>(), da.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX, da};
+}
+
+torch::Tensor adjT_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adjT_forward_kernel", ([&] {
+                                          adjT_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                 a.data_ptr<scalar_t>(),
+                                                                                 b.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return b;
+}
+
+std::vector<torch::Tensor> adjT_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adjT_backward_kernel", ([&] {
+                                          adjT_backward_kernel<group_t, scalar_t>(
+                                              grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(),
+                                              dX.data_ptr<scalar_t>(), da.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX, da};
+}
+
+torch::Tensor act_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor q = torch::zeros(p.sizes(), p.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act_forward_kernel", ([&] {
+                                          act_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                p.data_ptr<scalar_t>(),
+                                                                                q.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return q;
+}
+
+std::vector<torch::Tensor> act_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act_backward_kernel", ([&] {
+                                          act_backward_kernel<group_t, scalar_t>(
+                                              grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(),
+                                              dX.data_ptr<scalar_t>(), dp.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX, dp};
+}
+
+torch::Tensor act4_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor q = torch::zeros(p.sizes(), p.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act4_forward_kernel", ([&] {
+                                          act4_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                 p.data_ptr<scalar_t>(),
+                                                                                 q.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return q;
+}
+
+std::vector<torch::Tensor> act4_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act4_backward_kernel", ([&] {
+                                          act4_backward_kernel<group_t, scalar_t>(
+                                              grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(),
+                                              dX.data_ptr<scalar_t>(), dp.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX, dp};
+}
+
+torch::Tensor as_matrix_forward_cpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor T4x4 = torch::zeros({X.size(0), 4, 4}, X.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "as_matrix_forward_kernel", ([&] {
+                                          as_matrix_forward_kernel<group_t, scalar_t>(
+                                              X.data_ptr<scalar_t>(), T4x4.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return T4x4;
+}
+
+torch::Tensor orthogonal_projector_cpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor P;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "orthogonal_projector_kernel", ([&] {
+                                          P = torch::zeros({X.size(0), group_t::N, group_t::N}, X.options());
+                                          orthogonal_projector_kernel<group_t, scalar_t>(
+                                              X.data_ptr<scalar_t>(), P.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return P;
+}
+
+torch::Tensor jleft_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "jleft_forward_kernel", ([&] {
+                                          jleft_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                  a.data_ptr<scalar_t>(),
+                                                                                  b.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return b;
+}

third_party/vipe/csrc/lietorch_ext/lietorch_cpu.h

@@ -0,0 +1,50 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef LIETORCH_CPU_H_
+#define LIETORCH_CPU_H_
+
+#include <vector>
+#include <torch/extension.h>
+
+// unary operations
+torch::Tensor exp_forward_cpu(int, torch::Tensor);
+std::vector<torch::Tensor> exp_backward_cpu(int, torch::Tensor, torch::Tensor);
+
+torch::Tensor log_forward_cpu(int, torch::Tensor);
+std::vector<torch::Tensor> log_backward_cpu(int, torch::Tensor, torch::Tensor);
+
+torch::Tensor inv_forward_cpu(int, torch::Tensor);
+std::vector<torch::Tensor> inv_backward_cpu(int, torch::Tensor, torch::Tensor);
+
+// binary operations
+torch::Tensor mul_forward_cpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> mul_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor adj_forward_cpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> adj_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor adjT_forward_cpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> adjT_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor act_forward_cpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> act_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor act4_forward_cpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> act4_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+// conversion operations
+// std::vector<torch::Tensor> to_vec_backward_cpu(int, torch::Tensor, torch::Tensor);
+// std::vector<torch::Tensor> from_vec_backward_cpu(int, torch::Tensor, torch::Tensor);
+
+// utility operations
+torch::Tensor orthogonal_projector_cpu(int, torch::Tensor);
+
+torch::Tensor as_matrix_forward_cpu(int, torch::Tensor);
+
+torch::Tensor jleft_forward_cpu(int, torch::Tensor, torch::Tensor);
+
+#endif

third_party/vipe/csrc/lietorch_ext/lietorch_gpu.cu

@@ -0,0 +1,562 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include "lietorch_gpu.h"
+#include <eigen3/Eigen/Dense>
+
+#include "common.h"
+#include "dispatch.h"
+
+#include "so3.h"
+#include "rxso3.h"
+#include "se3.h"
+#include "sim3.h"
+
+#define GPU_1D_KERNEL_LOOP(i, n) \
+    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
+
+#define NUM_THREADS 256
+#define NUM_BLOCKS(batch_size) ((batch_size + NUM_THREADS - 1) / NUM_THREADS)
+
+template <typename Group, typename scalar_t>
+__global__ void exp_forward_kernel(const scalar_t* a_ptr, scalar_t* X_ptr, int num_threads) {
+    // exponential map forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Tangent a(a_ptr + i * Group::K);
+        Eigen::Map<Data>(X_ptr + i * Group::N) = Group::Exp(a).data();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void exp_backward_kernel(const scalar_t* grad, const scalar_t* a_ptr, scalar_t* da, int num_threads) {
+    // exponential map backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Tangent a(a_ptr + i * Group::K);
+        Grad dX(grad + i * Group::N);
+        Eigen::Map<Grad>(da + i * Group::K) = dX * Group::left_jacobian(a);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void log_forward_kernel(const scalar_t* X_ptr, scalar_t* a_ptr, int num_threads) {
+    // logarithm map forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Tangent a = Group(X_ptr + i * Group::N).Log();
+        Eigen::Map<Tangent>(a_ptr + i * Group::K) = a;
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void log_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int num_threads) {
+    // logarithm map backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Tangent a = Group(X_ptr + i * Group::N).Log();
+        Grad da(grad + i * Group::K);
+        Eigen::Map<Grad>(dX + i * Group::N) = da * Group::left_jacobian_inverse(a);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void inv_forward_kernel(const scalar_t* X_ptr, scalar_t* Y_ptr, int num_threads) {
+    // group inverse forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Eigen::Map<Data>(Y_ptr + i * Group::N) = X.inv().data();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void inv_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int num_threads) {
+    // group inverse backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group Y = Group(X_ptr + i * Group::N).inv();
+        Grad dY(grad + i * Group::N);
+        Eigen::Map<Grad>(dX + i * Group::N) = -dY * Y.Adj();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void mul_forward_kernel(const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* Z_ptr, int num_threads) {
+    // group multiplication forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group Z = Group(X_ptr + i * Group::N) * Group(Y_ptr + i * Group::N);
+        Eigen::Map<Data>(Z_ptr + i * Group::N) = Z.data();
+    }
+}
+
+template <class Group, typename scalar_t>
+__global__ void mul_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* dX,
+                                    scalar_t* dY, int num_threads) {
+    // group multiplication backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Grad dZ(grad + i * Group::N);
+        Group X(X_ptr + i * Group::N);
+        Eigen::Map<Grad>(dX + i * Group::N) = dZ;
+        Eigen::Map<Grad>(dY + i * Group::N) = dZ * X.Adj();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void adj_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int num_threads) {
+    // adjoint forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Tangent a(a_ptr + i * Group::K);
+        Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.Adj(a);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void adj_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX,
+                                    scalar_t* da, int num_threads) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Grad db(grad + i * Group::K);
+
+        Tangent a(a_ptr + i * Group::K);
+        Tangent b = X.Adj() * a;
+
+        Eigen::Map<Grad>(da + i * Group::K) = db * X.Adj();
+        Eigen::Map<Grad>(dX + i * Group::N) = -db * Group::adj(b);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void adjT_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int num_threads) {
+    // adjoint forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Tangent a(a_ptr + i * Group::K);
+        Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.AdjT(a);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void adjT_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX,
+                                     scalar_t* da, int num_threads) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Tangent db(grad + i * Group::K);
+        Grad a(a_ptr + i * Group::K);
+
+        Eigen::Map<Tangent>(da + i * Group::K) = X.Adj(db);
+        Eigen::Map<Grad>(dX + i * Group::N) = -a * Group::adj(X.Adj(db));
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void act_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int num_threads) {
+    // action on point forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Point = Eigen::Matrix<scalar_t, 3, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Point p(p_ptr + i * 3);
+        Eigen::Map<Point>(q_ptr + i * 3) = X * p;
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void act_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX,
+                                    scalar_t* dp, int num_threads) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Point = Eigen::Matrix<scalar_t, 3, 1>;
+    using PointGrad = Eigen::Matrix<scalar_t, 1, 3>;
+    using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Point p(p_ptr + i * 3);
+        PointGrad dq(grad + i * 3);
+
+        Eigen::Map<PointGrad>(dp + i * 3) = dq * X.Matrix4x4().block<3, 3>(0, 0);
+        Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act_jacobian(X * p);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void act4_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int num_threads) {
+    // action on point forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Point = Eigen::Matrix<scalar_t, 4, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Point p(p_ptr + i * 4);
+        Eigen::Map<Point>(q_ptr + i * 4) = X.act4(p);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void act4_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX,
+                                     scalar_t* dp, int num_threads) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Point = Eigen::Matrix<scalar_t, 4, 1>;
+    using PointGrad = Eigen::Matrix<scalar_t, 1, 4>;
+    using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Point p(p_ptr + i * 4);
+        PointGrad dq(grad + i * 4);
+
+        Eigen::Map<PointGrad>(dp + i * 4) = dq * X.Matrix4x4();
+        const Point q = X.act4(p);
+        Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act4_jacobian(q);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void as_matrix_forward_kernel(const scalar_t* X_ptr, scalar_t* T_ptr, int num_threads) {
+    // convert to 4x4 matrix representation
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Matrix4 = Eigen::Matrix<scalar_t, 4, 4, Eigen::RowMajor>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Eigen::Map<Matrix4>(T_ptr + i * 16) = X.Matrix4x4();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void orthogonal_projector_kernel(const scalar_t* X_ptr, scalar_t* P_ptr, int num_threads) {
+    // orthogonal projection matrix
+    using Proj = Eigen::Matrix<scalar_t, Group::N, Group::N, Eigen::RowMajor>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Eigen::Map<Proj>(P_ptr + i * Group::N * Group::N) = X.orthogonal_projector();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void jleft_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int num_threads) {
+    // left jacobian inverse action
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Tangent a(a_ptr + i * Group::K);
+        Tangent b = Group::left_jacobian_inverse(X.Log()) * a;
+        Eigen::Map<Tangent>(b_ptr + i * Group::K) = b;
+    }
+}
+
+// unary operations
+
+torch::Tensor exp_forward_gpu(int group_id, torch::Tensor a) {
+    int batch_size = a.size(0);
+    torch::Tensor X;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_forward_kernel", ([&] {
+                                          X = torch::zeros({batch_size, group_t::N}, a.options());
+                                          exp_forward_kernel<group_t, scalar_t>
+                                              <<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                                                  a.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return X;
+}
+
+std::vector<torch::Tensor> exp_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor a) {
+    int batch_size = a.size(0);
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, a.scalar_type(), "exp_backward_kernel", ([&] {
+            exp_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), da.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {da};
+}
+
+torch::Tensor log_forward_gpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor a;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_forward_kernel", ([&] {
+                                          a = torch::zeros({batch_size, group_t::K}, X.options());
+                                          log_forward_kernel<group_t, scalar_t>
+                                              <<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                                                  X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return a;
+}
+
+std::vector<torch::Tensor> log_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "log_backward_kernel", ([&] {
+            log_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX};
+}
+
+torch::Tensor inv_forward_gpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor Y = torch::zeros_like(X);
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "inv_forward_kernel", ([&] {
+                                          inv_forward_kernel<group_t, scalar_t>
+                                              <<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                                                  X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return Y;
+}
+
+std::vector<torch::Tensor> inv_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "inv_backward_kernel", ([&] {
+            inv_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX};
+}
+
+// binary operations
+torch::Tensor mul_forward_gpu(int group_id, torch::Tensor X, torch::Tensor Y) {
+    int batch_size = X.size(0);
+    torch::Tensor Z = torch::zeros_like(X);
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "mul_forward_kernel", ([&] {
+            mul_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), Z.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return Z;
+}
+
+std::vector<torch::Tensor> mul_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dY = torch::zeros(Y.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "mul_backward_kernel", ([&] {
+            mul_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
+                dY.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX, dY};
+}
+
+torch::Tensor adj_forward_gpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "adj_forward_kernel", ([&] {
+            adj_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), b.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return b;
+}
+
+std::vector<torch::Tensor> adj_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "adj_backward_kernel", ([&] {
+            adj_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
+                da.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX, da};
+}
+
+torch::Tensor adjT_forward_gpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "adjT_forward_kernel", ([&] {
+            adjT_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), b.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return b;
+}
+
+std::vector<torch::Tensor> adjT_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "adjT_backward_kernel", ([&] {
+            adjT_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
+                da.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX, da};
+}
+
+torch::Tensor act_forward_gpu(int group_id, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor q = torch::zeros(p.sizes(), p.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "act_forward_kernel", ([&] {
+            act_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), q.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return q;
+}
+
+std::vector<torch::Tensor> act_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "act_backward_kernel", ([&] {
+            act_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
+                dp.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX, dp};
+}
+
+torch::Tensor act4_forward_gpu(int group_id, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor q = torch::zeros(p.sizes(), p.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "act4_forward_kernel", ([&] {
+            act4_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), q.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return q;
+}
+
+std::vector<torch::Tensor> act4_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "act4_backward_kernel", ([&] {
+            act4_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
+                dp.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX, dp};
+}
+
+torch::Tensor as_matrix_forward_gpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor T4x4 = torch::zeros({X.size(0), 4, 4}, X.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "as_matrix_forward_kernel", ([&] {
+                                          as_matrix_forward_kernel<group_t, scalar_t>
+                                              <<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                                                  X.data_ptr<scalar_t>(), T4x4.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return T4x4;
+}
+
+torch::Tensor orthogonal_projector_gpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor P;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "orthogonal_projector_kernel", ([&] {
+                                          P = torch::zeros({X.size(0), group_t::N, group_t::N}, X.options());
+                                          orthogonal_projector_kernel<group_t, scalar_t>
+                                              <<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                                                  X.data_ptr<scalar_t>(), P.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return P;
+}
+
+torch::Tensor jleft_forward_gpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "jleft_forward_kernel", ([&] {
+            jleft_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), b.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return b;
+}

third_party/vipe/csrc/lietorch_ext/lietorch_gpu.h

@@ -0,0 +1,52 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef LIETORCH_GPU_H_
+#define LIETORCH_GPU_H_
+
+#include <vector>
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+// unary operations
+torch::Tensor exp_forward_gpu(int, torch::Tensor);
+std::vector<torch::Tensor> exp_backward_gpu(int, torch::Tensor, torch::Tensor);
+
+torch::Tensor log_forward_gpu(int, torch::Tensor);
+std::vector<torch::Tensor> log_backward_gpu(int, torch::Tensor, torch::Tensor);
+
+torch::Tensor inv_forward_gpu(int, torch::Tensor);
+std::vector<torch::Tensor> inv_backward_gpu(int, torch::Tensor, torch::Tensor);
+
+// binary operations
+torch::Tensor mul_forward_gpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> mul_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor adj_forward_gpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> adj_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor adjT_forward_gpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> adjT_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor act_forward_gpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> act_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor act4_forward_gpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> act4_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+// conversion operations
+// std::vector<torch::Tensor> to_vec_backward_gpu(int, torch::Tensor, torch::Tensor);
+// std::vector<torch::Tensor> from_vec_backward_gpu(int, torch::Tensor, torch::Tensor);
+
+// utility operators
+torch::Tensor orthogonal_projector_gpu(int, torch::Tensor);
+
+torch::Tensor as_matrix_forward_gpu(int, torch::Tensor);
+
+torch::Tensor jleft_forward_gpu(int, torch::Tensor, torch::Tensor);
+
+#endif

third_party/vipe/csrc/lietorch_ext/rxso3.h

@@ -0,0 +1,313 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef RxSO3_HEADER
+#define RxSO3_HEADER
+
+#include <stdio.h>
+#include <eigen3/Eigen/Dense>
+#include <eigen3/Eigen/Geometry>
+
+#include "common.h"
+
+template <typename Scalar>
+class RxSO3 {
+   public:
+    const static int constexpr K = 4;  // manifold dimension
+    const static int constexpr N = 5;  // embedding dimension
+
+    using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
+    using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
+
+    using Tangent = Eigen::Matrix<Scalar, K, 1>;
+    using Data = Eigen::Matrix<Scalar, N, 1>;
+
+    using Point = Eigen::Matrix<Scalar, 3, 1>;
+    using Point4 = Eigen::Matrix<Scalar, 4, 1>;
+
+    using Quaternion = Eigen::Quaternion<Scalar>;
+    using Transformation = Eigen::Matrix<Scalar, 3, 3>;
+    using Adjoint = Eigen::Matrix<Scalar, 4, 4>;
+
+    EIGEN_DEVICE_FUNC RxSO3(Quaternion const& q, Scalar const s) : unit_quaternion(q), scale(s) {
+        unit_quaternion.normalize();
+    };
+
+    EIGEN_DEVICE_FUNC RxSO3(const Scalar* data) : unit_quaternion(data), scale(data[4]) {
+        unit_quaternion.normalize();
+    };
+
+    EIGEN_DEVICE_FUNC RxSO3() {
+        unit_quaternion = Quaternion::Identity();
+        scale = Scalar(1.0);
+    }
+
+    EIGEN_DEVICE_FUNC RxSO3<Scalar> inv() { return RxSO3<Scalar>(unit_quaternion.conjugate(), 1.0 / scale); }
+
+    EIGEN_DEVICE_FUNC Data data() const {
+        Data data_vec;
+        data_vec << unit_quaternion.coeffs(), scale;
+        return data_vec;
+    }
+
+    EIGEN_DEVICE_FUNC RxSO3<Scalar> operator*(RxSO3<Scalar> const& other) {
+        return RxSO3<Scalar>(unit_quaternion * other.unit_quaternion, scale * other.scale);
+    }
+
+    EIGEN_DEVICE_FUNC Point operator*(Point const& p) const {
+        const Quaternion& q = unit_quaternion;
+        Point uv = q.vec().cross(p);
+        uv += uv;
+        return scale * (p + q.w() * uv + q.vec().cross(uv));
+    }
+
+    EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
+        Point4 p1;
+        p1 << this->operator*(p.template segment<3>(0)), p(3);
+        return p1;
+    }
+
+    EIGEN_DEVICE_FUNC Adjoint Adj() const {
+        Adjoint Ad = Adjoint::Identity();
+        Ad.template block<3, 3>(0, 0) = unit_quaternion.toRotationMatrix();
+        return Ad;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix() const { return scale * unit_quaternion.toRotationMatrix(); }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 4, 4> Matrix4x4() const {
+        Eigen::Matrix<Scalar, 4, 4> T;
+        T = Eigen::Matrix<Scalar, 4, 4>::Identity();
+        T.template block<3, 3>(0, 0) = Matrix();
+        return T;
+    }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 5, 5> orthogonal_projector() const {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 5, 5> J = Eigen::Matrix<Scalar, 5, 5>::Zero();
+
+        J.template block<3, 3>(0, 0) =
+            0.5 * (unit_quaternion.w() * Matrix3::Identity() + SO3<Scalar>::hat(-unit_quaternion.vec()));
+
+        J.template block<1, 3>(3, 0) = 0.5 * (-unit_quaternion.vec());
+
+        // scale
+        J(4, 3) = scale;
+
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Rotation() const { return unit_quaternion.toRotationMatrix(); }
+
+    EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
+
+    EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
+
+    EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& phi_sigma) {
+        Vector3 const phi = phi_sigma.template segment<3>(0);
+        return SO3<Scalar>::hat(phi) + phi(3) * Transformation::Identity();
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& phi_sigma) {
+        Vector3 const phi = phi_sigma.template segment<3>(0);
+        Matrix3 const Phi = SO3<Scalar>::hat(phi);
+
+        Adjoint ad = Adjoint::Zero();
+        ad.template block<3, 3>(0, 0) = Phi;
+
+        return ad;
+    }
+
+    EIGEN_DEVICE_FUNC Tangent Log() const {
+        using std::abs;
+        using std::atan;
+        using std::sqrt;
+
+        Scalar squared_n = unit_quaternion.vec().squaredNorm();
+        Scalar w = unit_quaternion.w();
+        Scalar two_atan_nbyw_by_n;
+
+        /// Atan-based log thanks to
+        ///
+        /// C. Hertzberg et al.:
+        /// "Integrating Generic Sensor Fusion Algorithms with Sound State
+        /// Representation through Encapsulation of Manifolds"
+        /// Information Fusion, 2011
+
+        if (squared_n < EPS * EPS) {
+            two_atan_nbyw_by_n = Scalar(2) / w - Scalar(2.0 / 3.0) * (squared_n) / (w * w * w);
+        } else {
+            Scalar n = sqrt(squared_n);
+            if (abs(w) < EPS) {
+                if (w > Scalar(0)) {
+                    two_atan_nbyw_by_n = PI / n;
+                } else {
+                    two_atan_nbyw_by_n = -PI / n;
+                }
+            } else {
+                two_atan_nbyw_by_n = Scalar(2) * atan(n / w) / n;
+            }
+        }
+
+        Tangent phi_sigma;
+        phi_sigma << two_atan_nbyw_by_n * unit_quaternion.vec(), log(scale);
+
+        return phi_sigma;
+    }
+
+    EIGEN_DEVICE_FUNC static RxSO3<Scalar> Exp(Tangent const& phi_sigma) {
+        Vector3 phi = phi_sigma.template segment<3>(0);
+        Scalar scale = exp(phi_sigma(3));
+
+        Scalar theta2 = phi.squaredNorm();
+        Scalar theta = sqrt(theta2);
+        Scalar imag_factor;
+        Scalar real_factor;
+
+        if (theta < EPS) {
+            Scalar theta4 = theta2 * theta2;
+            imag_factor = Scalar(0.5) - Scalar(1.0 / 48.0) * theta2 + Scalar(1.0 / 3840.0) * theta4;
+            real_factor = Scalar(1) - Scalar(1.0 / 8.0) * theta2 + Scalar(1.0 / 384.0) * theta4;
+        } else {
+            imag_factor = sin(.5 * theta) / theta;
+            real_factor = cos(.5 * theta);
+        }
+
+        Quaternion q(real_factor, imag_factor * phi.x(), imag_factor * phi.y(), imag_factor * phi.z());
+        return RxSO3<Scalar>(q, scale);
+    }
+
+    EIGEN_DEVICE_FUNC static Matrix3 calcW(Tangent const& phi_sigma) {
+        // left jacobian
+        using std::abs;
+        Matrix3 const I = Matrix3::Identity();
+        Scalar const one(1);
+        Scalar const half(0.5);
+
+        Vector3 const phi = phi_sigma.template segment<3>(0);
+        Scalar const sigma = phi_sigma(3);
+        Scalar const theta = phi.norm();
+
+        Matrix3 const Phi = SO3<Scalar>::hat(phi);
+        Matrix3 const Phi2 = Phi * Phi;
+        Scalar const scale = exp(sigma);
+
+        Scalar A, B, C;
+        if (abs(sigma) < EPS) {
+            C = one;
+            if (abs(theta) < EPS) {
+                A = half;
+                B = Scalar(1. / 6.);
+            } else {
+                Scalar theta_sq = theta * theta;
+                A = (one - cos(theta)) / theta_sq;
+                B = (theta - sin(theta)) / (theta_sq * theta);
+            }
+        } else {
+            C = (scale - one) / sigma;
+            if (abs(theta) < EPS) {
+                Scalar sigma_sq = sigma * sigma;
+                A = ((sigma - one) * scale + one) / sigma_sq;
+                B = (scale * half * sigma_sq + scale - one - sigma * scale) / (sigma_sq * sigma);
+            } else {
+                Scalar theta_sq = theta * theta;
+                Scalar a = scale * sin(theta);
+                Scalar b = scale * cos(theta);
+                Scalar c = theta_sq + sigma * sigma;
+                A = (a * sigma + (one - b) * theta) / (theta * c);
+                B = (C - ((b - one) * sigma + a * theta) / (c)) * one / (theta_sq);
+            }
+        }
+        return A * Phi + B * Phi2 + C * I;
+    }
+
+    EIGEN_DEVICE_FUNC static Matrix3 calcWInv(Tangent const& phi_sigma) {
+        // left jacobian inverse
+        Matrix3 const I = Matrix3::Identity();
+        Scalar const half(0.5);
+        Scalar const one(1);
+        Scalar const two(2);
+
+        Vector3 const phi = phi_sigma.template segment<3>(0);
+        Scalar const sigma = phi_sigma(3);
+        Scalar const theta = phi.norm();
+        Scalar const scale = exp(sigma);
+
+        Matrix3 const Phi = SO3<Scalar>::hat(phi);
+        Matrix3 const Phi2 = Phi * Phi;
+        Scalar const scale_sq = scale * scale;
+        Scalar const theta_sq = theta * theta;
+        Scalar const sin_theta = sin(theta);
+        Scalar const cos_theta = cos(theta);
+
+        Scalar a, b, c;
+        if (abs(sigma * sigma) < EPS) {
+            c = one - half * sigma;
+            a = -half;
+            if (abs(theta_sq) < EPS) {
+                b = Scalar(1. / 12.);
+            } else {
+                b = (theta * sin_theta + two * cos_theta - two) / (two * theta_sq * (cos_theta - one));
+            }
+        } else {
+            Scalar const scale_cu = scale_sq * scale;
+            c = sigma / (scale - one);
+            if (abs(theta_sq) < EPS) {
+                a = (-sigma * scale + scale - one) / ((scale - one) * (scale - one));
+                b = (scale_sq * sigma - two * scale_sq + scale * sigma + two * scale) /
+                    (two * scale_cu - Scalar(6) * scale_sq + Scalar(6) * scale - two);
+            } else {
+                Scalar const s_sin_theta = scale * sin_theta;
+                Scalar const s_cos_theta = scale * cos_theta;
+                a = (theta * s_cos_theta - theta - sigma * s_sin_theta) /
+                    (theta * (scale_sq - two * s_cos_theta + one));
+                b = -scale *
+                    (theta * s_sin_theta - theta * sin_theta + sigma * s_cos_theta - scale * sigma + sigma * cos_theta -
+                     sigma) /
+                    (theta_sq * (scale_cu - two * scale * s_cos_theta - scale_sq + two * s_cos_theta + scale - one));
+            }
+        }
+        return a * Phi + b * Phi2 + c * I;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& phi_sigma) {
+        // left jacobian
+        Adjoint J = Adjoint::Identity();
+        Vector3 phi = phi_sigma.template segment<3>(0);
+        J.template block<3, 3>(0, 0) = SO3<Scalar>::left_jacobian(phi);
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& phi_sigma) {
+        // left jacobian inverse
+        Adjoint Jinv = Adjoint::Identity();
+        Vector3 phi = phi_sigma.template segment<3>(0);
+        Jinv.template block<3, 3>(0, 0) = SO3<Scalar>::left_jacobian_inverse(phi);
+        return Jinv;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 4> act_jacobian(Point const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 3, 4> Ja;
+        Ja << SO3<Scalar>::hat(-p), p;
+        return Ja;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 4> act4_jacobian(Point4 const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 4, 4> J = Eigen::Matrix<Scalar, 4, 4>::Zero();
+        J.template block<3, 3>(0, 0) = SO3<Scalar>::hat(-p.template segment<3>(0));
+        J.template block<3, 1>(0, 3) = p.template segment<3>(0);
+        return J;
+    }
+
+   private:
+    Quaternion unit_quaternion;
+    Scalar scale;
+};
+
+#endif

third_party/vipe/csrc/lietorch_ext/se3.h

@@ -0,0 +1,214 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef SE3_HEADER
+#define SE3_HEADER
+
+#include <stdio.h>
+#include <eigen3/Eigen/Dense>
+#include <eigen3/Eigen/Geometry>
+
+#include "common.h"
+#include "so3.h"
+
+template <typename Scalar>
+class SE3 {
+   public:
+    const static int constexpr K = 6;  // manifold dimension
+    const static int constexpr N = 7;  // embedding dimension
+
+    using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
+    using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
+
+    using Tangent = Eigen::Matrix<Scalar, K, 1>;
+    using Point = Eigen::Matrix<Scalar, 3, 1>;
+    using Point4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Data = Eigen::Matrix<Scalar, N, 1>;
+    using Transformation = Eigen::Matrix<Scalar, 4, 4>;
+    using Adjoint = Eigen::Matrix<Scalar, K, K>;
+
+    EIGEN_DEVICE_FUNC SE3() { translation = Vector3::Zero(); }
+
+    EIGEN_DEVICE_FUNC SE3(SO3<Scalar> const& so3, Vector3 const& t) : so3(so3), translation(t) {};
+
+    EIGEN_DEVICE_FUNC SE3(const Scalar* data) : translation(data), so3(data + 3) {};
+
+    EIGEN_DEVICE_FUNC SE3<Scalar> inv() { return SE3(so3.inv(), -(so3.inv() * translation)); }
+
+    EIGEN_DEVICE_FUNC Data data() const {
+        Data data_vec;
+        data_vec << translation, so3.data();
+        return data_vec;
+    }
+
+    EIGEN_DEVICE_FUNC SE3<Scalar> operator*(SE3<Scalar> const& other) {
+        return SE3(so3 * other.so3, translation + so3 * other.translation);
+    }
+
+    EIGEN_DEVICE_FUNC Point operator*(Point const& p) const { return so3 * p + translation; }
+
+    EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
+        Point4 p1;
+        p1 << so3 * p.template segment<3>(0) + translation * p(3), p(3);
+        return p1;
+    }
+
+    EIGEN_DEVICE_FUNC Adjoint Adj() const {
+        Matrix3 R = so3.Matrix();
+        Matrix3 tx = SO3<Scalar>::hat(translation);
+        Matrix3 Zer = Matrix3::Zero();
+
+        Adjoint Ad;
+        Ad << R, tx * R, Zer, R;
+
+        return Ad;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix() const {
+        Transformation T = Transformation::Identity();
+        T.template block<3, 3>(0, 0) = so3.Matrix();
+        T.template block<3, 1>(0, 3) = translation;
+        return T;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix4x4() const { return Matrix(); }
+
+    EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
+
+    EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
+
+    EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& tau_phi) {
+        Vector3 tau = tau_phi.template segment<3>(0);
+        Vector3 phi = tau_phi.template segment<3>(3);
+
+        Transformation TauPhi = Transformation::Zero();
+        TauPhi.template block<3, 3>(0, 0) = SO3<Scalar>::hat(phi);
+        TauPhi.template block<3, 1>(0, 3) = tau;
+
+        return TauPhi;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& tau_phi) {
+        Vector3 tau = tau_phi.template segment<3>(0);
+        Vector3 phi = tau_phi.template segment<3>(3);
+
+        Matrix3 Tau = SO3<Scalar>::hat(tau);
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+        Matrix3 Zer = Matrix3::Zero();
+
+        Adjoint ad;
+        ad << Phi, Tau, Zer, Phi;
+
+        return ad;
+    }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 7, 7> orthogonal_projector() const {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 7, 7> J = Eigen::Matrix<Scalar, 7, 7>::Zero();
+        J.template block<3, 3>(0, 0) = Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-translation);
+        J.template block<4, 4>(3, 3) = so3.orthogonal_projector();
+
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC Tangent Log() const {
+        Vector3 phi = so3.Log();
+        Matrix3 Vinv = SO3<Scalar>::left_jacobian_inverse(phi);
+
+        Tangent tau_phi;
+        tau_phi << Vinv * translation, phi;
+
+        return tau_phi;
+    }
+
+    EIGEN_DEVICE_FUNC static SE3<Scalar> Exp(Tangent const& tau_phi) {
+        Vector3 tau = tau_phi.template segment<3>(0);
+        Vector3 phi = tau_phi.template segment<3>(3);
+
+        SO3<Scalar> so3 = SO3<Scalar>::Exp(phi);
+        Vector3 t = SO3<Scalar>::left_jacobian(phi) * tau;
+
+        return SE3<Scalar>(so3, t);
+    }
+
+    EIGEN_DEVICE_FUNC static Matrix3 calcQ(Tangent const& tau_phi) {
+        // Q matrix
+        Vector3 tau = tau_phi.template segment<3>(0);
+        Vector3 phi = tau_phi.template segment<3>(3);
+        Matrix3 Tau = SO3<Scalar>::hat(tau);
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+
+        Scalar theta = phi.norm();
+        Scalar theta_pow2 = theta * theta;
+        Scalar theta_pow4 = theta_pow2 * theta_pow2;
+
+        Scalar coef1 = (theta < EPS) ? Scalar(1.0 / 6.0) - Scalar(1.0 / 120.0) * theta_pow2
+                                     : (theta - sin(theta)) / (theta_pow2 * theta);
+
+        Scalar coef2 = (theta < EPS) ? Scalar(1.0 / 24.0) - Scalar(1.0 / 720.0) * theta_pow2
+                                     : (theta_pow2 + 2 * cos(theta) - 2) / (2 * theta_pow4);
+
+        Scalar coef3 = (theta < EPS) ? Scalar(1.0 / 120.0) - Scalar(1.0 / 2520.0) * theta_pow2
+                                     : (2 * theta - 3 * sin(theta) + theta * cos(theta)) / (2 * theta_pow4 * theta);
+
+        Matrix3 Q = Scalar(0.5) * Tau + coef1 * (Phi * Tau + Tau * Phi + Phi * Tau * Phi) +
+                    coef2 * (Phi * Phi * Tau + Tau * Phi * Phi - 3 * Phi * Tau * Phi) +
+                    coef3 * (Phi * Tau * Phi * Phi + Phi * Phi * Tau * Phi);
+
+        return Q;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& tau_phi) {
+        // left jacobian
+        Vector3 phi = tau_phi.template segment<3>(3);
+        Matrix3 J = SO3<Scalar>::left_jacobian(phi);
+        Matrix3 Q = SE3<Scalar>::calcQ(tau_phi);
+        Matrix3 Zer = Matrix3::Zero();
+
+        Adjoint J6x6;
+        J6x6 << J, Q, Zer, J;
+
+        return J6x6;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& tau_phi) {
+        // left jacobian inverse
+        Vector3 tau = tau_phi.template segment<3>(0);
+        Vector3 phi = tau_phi.template segment<3>(3);
+        Matrix3 Jinv = SO3<Scalar>::left_jacobian_inverse(phi);
+        Matrix3 Q = SE3<Scalar>::calcQ(tau_phi);
+        Matrix3 Zer = Matrix3::Zero();
+
+        Adjoint J6x6;
+        J6x6 << Jinv, -Jinv * Q * Jinv, Zer, Jinv;
+
+        return J6x6;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 6> act_jacobian(Point const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 3, 6> J;
+        J.template block<3, 3>(0, 0) = Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p);
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 6> act4_jacobian(Point4 const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 4, 6> J = Eigen::Matrix<Scalar, 4, 6>::Zero();
+        J.template block<3, 3>(0, 0) = p(3) * Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p.template segment<3>(0));
+        return J;
+    }
+
+   private:
+    SO3<Scalar> so3;
+    Vector3 translation;
+};
+
+#endif

third_party/vipe/csrc/lietorch_ext/sim3.h

@@ -0,0 +1,203 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef Sim3_HEADER
+#define Sim3_HEADER
+
+#include <stdio.h>
+#include <iostream>
+
+#include <eigen3/Eigen/Dense>
+#include <eigen3/Eigen/Geometry>
+
+#include "common.h"
+#include "so3.h"
+#include "rxso3.h"
+
+template <typename Scalar>
+class Sim3 {
+   public:
+    const static int constexpr K = 7;  // manifold dimension
+    const static int constexpr N = 8;  // embedding dimension
+
+    using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
+    using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
+
+    using Tangent = Eigen::Matrix<Scalar, K, 1>;
+    using Point = Eigen::Matrix<Scalar, 3, 1>;
+    using Point4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Data = Eigen::Matrix<Scalar, N, 1>;
+    using Transformation = Eigen::Matrix<Scalar, 4, 4>;
+    using Adjoint = Eigen::Matrix<Scalar, K, K>;
+
+    EIGEN_DEVICE_FUNC Sim3() { translation = Vector3::Zero(); }
+
+    EIGEN_DEVICE_FUNC Sim3(RxSO3<Scalar> const& rxso3, Vector3 const& t) : rxso3(rxso3), translation(t) {};
+
+    EIGEN_DEVICE_FUNC Sim3(const Scalar* data) : translation(data), rxso3(data + 3) {};
+
+    EIGEN_DEVICE_FUNC Sim3<Scalar> inv() { return Sim3<Scalar>(rxso3.inv(), -(rxso3.inv() * translation)); }
+
+    EIGEN_DEVICE_FUNC Data data() const {
+        Data data_vec;
+        data_vec << translation, rxso3.data();
+        return data_vec;
+    }
+
+    EIGEN_DEVICE_FUNC Sim3<Scalar> operator*(Sim3<Scalar> const& other) {
+        return Sim3(rxso3 * other.rxso3, translation + rxso3 * other.translation);
+    }
+
+    EIGEN_DEVICE_FUNC Point operator*(Point const& p) const { return (rxso3 * p) + translation; }
+
+    EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
+        Point4 p1;
+        p1 << rxso3 * p.template segment<3>(0) + p(3) * translation, p(3);
+        return p1;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix() const {
+        Transformation T = Transformation::Identity();
+        T.template block<3, 3>(0, 0) = rxso3.Matrix();
+        T.template block<3, 1>(0, 3) = translation;
+        return T;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix4x4() const {
+        Transformation T = Transformation::Identity();
+        T.template block<3, 3>(0, 0) = rxso3.Matrix();
+        T.template block<3, 1>(0, 3) = translation;
+        return T;
+    }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 8, 8> orthogonal_projector() const {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 8, 8> J = Eigen::Matrix<Scalar, 8, 8>::Zero();
+        J.template block<3, 3>(0, 0) = Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-translation);
+        J.template block<3, 1>(0, 6) = translation;
+        J.template block<5, 5>(3, 3) = rxso3.orthogonal_projector();
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC Adjoint Adj() const {
+        Adjoint Ad = Adjoint::Identity();
+        Matrix3 sR = rxso3.Matrix();
+        Matrix3 tx = SO3<Scalar>::hat(translation);
+        Matrix3 R = rxso3.Rotation();
+
+        Ad.template block<3, 3>(0, 0) = sR;
+        Ad.template block<3, 3>(0, 3) = tx * R;
+        Ad.template block<3, 1>(0, 6) = -translation;
+        Ad.template block<3, 3>(3, 3) = R;
+
+        return Ad;
+    }
+
+    EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
+
+    EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
+
+    EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& tau_phi_sigma) {
+        Vector3 tau = tau_phi_sigma.template segment<3>(0);
+        Vector3 phi = tau_phi_sigma.template segment<3>(3);
+        Scalar sigma = tau_phi_sigma(6);
+
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+        Matrix3 I = Matrix3::Identity();
+
+        Transformation Omega = Transformation::Zero();
+        Omega.template block<3, 3>(0, 0) = Phi + sigma * I;
+        Omega.template block<3, 1>(0, 3) = tau;
+
+        return Omega;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& tau_phi_sigma) {
+        Adjoint ad = Adjoint::Zero();
+        Vector3 tau = tau_phi_sigma.template segment<3>(0);
+        Vector3 phi = tau_phi_sigma.template segment<3>(3);
+        Scalar sigma = tau_phi_sigma(6);
+
+        Matrix3 Tau = SO3<Scalar>::hat(tau);
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+        Matrix3 I = Matrix3::Identity();
+
+        ad.template block<3, 3>(0, 0) = Phi + sigma * I;
+        ad.template block<3, 3>(0, 3) = Tau;
+        ad.template block<3, 1>(0, 6) = -tau;
+        ad.template block<3, 3>(3, 3) = Phi;
+
+        return ad;
+    }
+
+    EIGEN_DEVICE_FUNC Tangent Log() const {
+        // logarithm map
+        Vector4 phi_sigma = rxso3.Log();
+        Matrix3 W = RxSO3<Scalar>::calcW(phi_sigma);
+
+        Tangent tau_phi_sigma;
+        tau_phi_sigma << W.inverse() * translation, phi_sigma;
+
+        return tau_phi_sigma;
+    }
+
+    EIGEN_DEVICE_FUNC static Sim3<Scalar> Exp(Tangent const& tau_phi_sigma) {
+        // exponential map
+        Vector3 tau = tau_phi_sigma.template segment<3>(0);
+        Vector4 phi_sigma = tau_phi_sigma.template segment<4>(3);
+
+        RxSO3<Scalar> rxso3 = RxSO3<Scalar>::Exp(phi_sigma);
+        Matrix3 W = RxSO3<Scalar>::calcW(phi_sigma);
+
+        return Sim3<Scalar>(rxso3, W * tau);
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& tau_phi_sigma) {
+        // left jacobian
+        Adjoint const Xi = adj(tau_phi_sigma);
+        Adjoint const Xi2 = Xi * Xi;
+        Adjoint const Xi4 = Xi2 * Xi2;
+
+        return Adjoint::Identity() + Scalar(1.0 / 2.0) * Xi + Scalar(1.0 / 6.0) * Xi2 + Scalar(1.0 / 24.0) * Xi * Xi2 +
+               Scalar(1.0 / 120.0) * Xi4;
+        +Scalar(1.0 / 720.0) * Xi* Xi4;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& tau_phi_sigma) {
+        // left jacobian inverse
+        Adjoint const Xi = adj(tau_phi_sigma);
+        Adjoint const Xi2 = Xi * Xi;
+        Adjoint const Xi4 = Xi2 * Xi2;
+
+        return Adjoint::Identity() - Scalar(1.0 / 2.0) * Xi + Scalar(1.0 / 12.0) * Xi2 - Scalar(1.0 / 720.0) * Xi4;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 7> act_jacobian(Point const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 3, 7> J;
+        J.template block<3, 3>(0, 0) = Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p);
+        J.template block<3, 1>(0, 6) = p;
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 7> act4_jacobian(Point4 const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 4, 7> J = Eigen::Matrix<Scalar, 4, 7>::Zero();
+        J.template block<3, 3>(0, 0) = p(3) * Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p.template segment<3>(0));
+        J.template block<3, 1>(0, 6) = p.template segment<3>(0);
+        return J;
+    }
+
+   private:
+    Vector3 translation;
+    RxSO3<Scalar> rxso3;
+};
+
+#endif

third_party/vipe/csrc/lietorch_ext/so3.h

@@ -0,0 +1,202 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef SO3_HEADER
+#define SO3_HEADER
+
+#include <cuda.h>
+#include <stdio.h>
+#include <eigen3/Eigen/Dense>
+#include <eigen3/Eigen/Geometry>
+
+#include "common.h"
+
+template <typename Scalar>
+class SO3 {
+   public:
+    const static int constexpr K = 3;  // manifold dimension
+    const static int constexpr N = 4;  // embedding dimension
+
+    using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
+    using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
+
+    using Tangent = Eigen::Matrix<Scalar, K, 1>;
+    using Data = Eigen::Matrix<Scalar, N, 1>;
+
+    using Point = Eigen::Matrix<Scalar, 3, 1>;
+    using Point4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Transformation = Eigen::Matrix<Scalar, 3, 3>;
+    using Adjoint = Eigen::Matrix<Scalar, K, K>;
+    using Quaternion = Eigen::Quaternion<Scalar>;
+
+    EIGEN_DEVICE_FUNC SO3(Quaternion const& q) : unit_quaternion(q) { unit_quaternion.normalize(); };
+
+    EIGEN_DEVICE_FUNC SO3(const Scalar* data) : unit_quaternion(data) { unit_quaternion.normalize(); };
+
+    EIGEN_DEVICE_FUNC SO3() { unit_quaternion = Quaternion::Identity(); }
+
+    EIGEN_DEVICE_FUNC SO3<Scalar> inv() { return SO3<Scalar>(unit_quaternion.conjugate()); }
+
+    EIGEN_DEVICE_FUNC Data data() const { return unit_quaternion.coeffs(); }
+
+    EIGEN_DEVICE_FUNC SO3<Scalar> operator*(SO3<Scalar> const& other) {
+        return SO3(unit_quaternion * other.unit_quaternion);
+    }
+
+    EIGEN_DEVICE_FUNC Point operator*(Point const& p) const {
+        const Quaternion& q = unit_quaternion;
+        Point uv = q.vec().cross(p);
+        uv += uv;
+        return p + q.w() * uv + q.vec().cross(uv);
+    }
+
+    EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
+        Point4 p1;
+        p1 << this->operator*(p.template segment<3>(0)), p(3);
+        return p1;
+    }
+
+    EIGEN_DEVICE_FUNC Adjoint Adj() const { return unit_quaternion.toRotationMatrix(); }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix() const { return unit_quaternion.toRotationMatrix(); }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 4, 4> Matrix4x4() const {
+        Eigen::Matrix<Scalar, 4, 4> T = Eigen::Matrix<Scalar, 4, 4>::Identity();
+        T.template block<3, 3>(0, 0) = Matrix();
+        return T;
+    }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 4, 4> orthogonal_projector() const {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 4, 4> J = Eigen::Matrix<Scalar, 4, 4>::Zero();
+        J.template block<3, 3>(0, 0) =
+            0.5 * (unit_quaternion.w() * Matrix3::Identity() + SO3<Scalar>::hat(-unit_quaternion.vec()));
+
+        J.template block<1, 3>(3, 0) = 0.5 * (-unit_quaternion.vec());
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
+
+    EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
+
+    EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& phi) {
+        Transformation Phi;
+        Phi << 0.0, -phi(2), phi(1), phi(2), 0.0, -phi(0), -phi(1), phi(0), 0.0;
+
+        return Phi;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& phi) { return SO3<Scalar>::hat(phi); }
+
+    EIGEN_DEVICE_FUNC Tangent Log() const {
+        using std::abs;
+        using std::atan;
+        using std::sqrt;
+        Scalar squared_n = unit_quaternion.vec().squaredNorm();
+        Scalar w = unit_quaternion.w();
+
+        Scalar two_atan_nbyw_by_n;
+
+        /// Atan-based log thanks to
+        ///
+        /// C. Hertzberg et al.:
+        /// "Integrating Generic Sensor Fusion Algorithms with Sound State
+        /// Representation through Encapsulation of Manifolds"
+        /// Information Fusion, 2011
+
+        if (squared_n < EPS * EPS) {
+            // If quaternion is normalized and n=0, then w should be 1;
+            // w=0 should never happen here!
+            Scalar squared_w = w * w;
+            two_atan_nbyw_by_n = Scalar(2) / w - Scalar(2.0 / 3.0) * (squared_n) / (w * squared_w);
+        } else {
+            Scalar n = sqrt(squared_n);
+            if (abs(w) < EPS) {
+                if (w > Scalar(0)) {
+                    two_atan_nbyw_by_n = Scalar(PI) / n;
+                } else {
+                    two_atan_nbyw_by_n = -Scalar(PI) / n;
+                }
+            } else {
+                two_atan_nbyw_by_n = Scalar(2) * atan(n / w) / n;
+            }
+        }
+
+        return two_atan_nbyw_by_n * unit_quaternion.vec();
+    }
+
+    EIGEN_DEVICE_FUNC static SO3<Scalar> Exp(Tangent const& phi) {
+        Scalar theta2 = phi.squaredNorm();
+        Scalar theta = sqrt(theta2);
+        Scalar imag_factor;
+        Scalar real_factor;
+
+        if (theta < EPS) {
+            Scalar theta4 = theta2 * theta2;
+            imag_factor = Scalar(0.5) - Scalar(1.0 / 48.0) * theta2 + Scalar(1.0 / 3840.0) * theta4;
+            real_factor = Scalar(1) - Scalar(1.0 / 8.0) * theta2 + Scalar(1.0 / 384.0) * theta4;
+        } else {
+            imag_factor = sin(.5 * theta) / theta;
+            real_factor = cos(.5 * theta);
+        }
+
+        Quaternion q(real_factor, imag_factor * phi.x(), imag_factor * phi.y(), imag_factor * phi.z());
+        return SO3<Scalar>(q);
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& phi) {
+        // left jacobian
+        Matrix3 I = Matrix3::Identity();
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+        Matrix3 Phi2 = Phi * Phi;
+
+        Scalar theta2 = phi.squaredNorm();
+        Scalar theta = sqrt(theta2);
+
+        Scalar coef1 = (theta < EPS) ? Scalar(1.0 / 2.0) - Scalar(1.0 / 24.0) * theta2 : (1.0 - cos(theta)) / theta2;
+
+        Scalar coef2 =
+            (theta < EPS) ? Scalar(1.0 / 6.0) - Scalar(1.0 / 120.0) * theta2 : (theta - sin(theta)) / (theta2 * theta);
+
+        return I + coef1 * Phi + coef2 * Phi2;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& phi) {
+        // left jacobian inverse
+        Matrix3 I = Matrix3::Identity();
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+        Matrix3 Phi2 = Phi * Phi;
+
+        Scalar theta2 = phi.squaredNorm();
+        Scalar theta = sqrt(theta2);
+        Scalar half_theta = Scalar(.5) * theta;
+
+        Scalar coef2 = (theta < EPS)
+                           ? Scalar(1.0 / 12.0)
+                           : (Scalar(1) - theta * cos(half_theta) / (Scalar(2) * sin(half_theta))) / (theta * theta);
+
+        return I + Scalar(-0.5) * Phi + coef2 * Phi2;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 3> act_jacobian(Point const& p) {
+        // jacobian action on a point
+        return SO3<Scalar>::hat(-p);
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 3> act4_jacobian(Point4 const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 4, 3> J = Eigen::Matrix<Scalar, 4, 3>::Zero();
+        J.template block<3, 3>(0, 0) = SO3<Scalar>::hat(-p.template segment<3>(0));
+        return J;
+    }
+
+   private:
+    Quaternion unit_quaternion;
+};
+
+#endif

third_party/vipe/csrc/scatter_ext/cuda/atomics.cuh

@@ -0,0 +1,300 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#pragma once
+
+#define ATOMIC(NAME)                                                                                                  \
+    template <typename scalar, size_t size>                                                                           \
+    struct Atomic##NAME##IntegerImpl;                                                                                 \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##IntegerImpl<scalar, 1> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            uint32_t *address_as_ui = (uint32_t *)(address - ((size_t)address & 3));                                  \
+            uint32_t old = *address_as_ui;                                                                            \
+            uint32_t shift = ((size_t)address & 3) * 8;                                                               \
+            uint32_t sum;                                                                                             \
+            uint32_t assumed;                                                                                         \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                sum = OP(val, scalar((old >> shift) & 0xff));                                                         \
+                old = (old & ~(0x000000ff << shift)) | (sum << shift);                                                \
+                old = atomicCAS(address_as_ui, assumed, old);                                                         \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##IntegerImpl<scalar, 2> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            uint32_t *address_as_ui = (uint32_t *)((char *)address - ((size_t)address & 2));                          \
+            uint32_t old = *address_as_ui;                                                                            \
+            uint32_t sum;                                                                                             \
+            uint32_t newval;                                                                                          \
+            uint32_t assumed;                                                                                         \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                sum = OP(val, (size_t)address & 2 ? scalar(old >> 16) : scalar(old & 0xffff));                        \
+                newval = (size_t)address & 2 ? (old & 0xffff) | (sum << 16) : (old & 0xffff0000) | sum;               \
+                old = atomicCAS(address_as_ui, assumed, newval);                                                      \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##IntegerImpl<scalar, 4> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            uint32_t *address_as_ui = (uint32_t *)address;                                                            \
+            uint32_t old = *address_as_ui;                                                                            \
+            uint32_t assumed;                                                                                         \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                old = atomicCAS(address_as_ui, assumed, OP(val, (scalar)old));                                        \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##IntegerImpl<scalar, 8> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            unsigned long long *address_as_ull = (unsigned long long *)address;                                       \
+            unsigned long long old = *address_as_ull;                                                                 \
+            unsigned long long assumed;                                                                               \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                old = atomicCAS(address_as_ull, assumed, OP(val, (scalar)old));                                       \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar, size_t size>                                                                           \
+    struct Atomic##NAME##DecimalImpl;                                                                                 \
+                                                                                                                      \
+    template <>                                                                                                       \
+    struct Atomic##NAME##DecimalImpl<at::Half, 2> {                                                                   \
+        inline __device__ void operator()(at::Half *address, at::Half val) {                                          \
+            unsigned int *address_as_ui = (unsigned int *)((char *)address - ((size_t)address & 2));                  \
+            unsigned int old = *address_as_ui;                                                                        \
+            unsigned int assumed;                                                                                     \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                at::Half hsum;                                                                                        \
+                hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff);                                          \
+                hsum = OP(hsum, val);                                                                                 \
+                old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16) : (old & 0xffff0000) | hsum.x;            \
+                old = atomicCAS(address_as_ui, assumed, old);                                                         \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <>                                                                                                       \
+    struct Atomic##NAME##DecimalImpl<at::BFloat16, 2> {                                                               \
+        inline __device__ void operator()(at::BFloat16 *address, at::BFloat16 val) {                                  \
+            unsigned int *address_as_ui = (unsigned int *)((char *)address - ((size_t)address & 2));                  \
+            unsigned int old = *address_as_ui;                                                                        \
+            unsigned int assumed;                                                                                     \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                at::BFloat16 hsum;                                                                                    \
+                hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff);                                          \
+                hsum = OP(hsum, val);                                                                                 \
+                old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16) : (old & 0xffff0000) | hsum.x;            \
+                old = atomicCAS(address_as_ui, assumed, old);                                                         \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##DecimalImpl<scalar, 4> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            int *address_as_i = (int *)address;                                                                       \
+            int old = *address_as_i;                                                                                  \
+            int assumed;                                                                                              \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                old = atomicCAS(address_as_i, assumed, __float_as_int(OP(val, __int_as_float(assumed))));             \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##DecimalImpl<scalar, 8> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            unsigned long long int *address_as_ull = (unsigned long long int *)address;                               \
+            unsigned long long int old = *address_as_ull;                                                             \
+            unsigned long long int assumed;                                                                           \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                old =                                                                                                 \
+                    atomicCAS(address_as_ull, assumed, __double_as_longlong(OP(val, __longlong_as_double(assumed)))); \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };
+
+#define OP(X, Y) Y + X
+ATOMIC(Add)
+#undef OP
+static inline __device__ void atomAdd(uint8_t *address, uint8_t val) {
+    AtomicAddIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
+}
+static inline __device__ void atomAdd(int8_t *address, int8_t val) {
+    AtomicAddIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
+}
+static inline __device__ void atomAdd(int16_t *address, int16_t val) {
+    AtomicAddIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
+}
+static inline __device__ void atomAdd(int32_t *address, int32_t val) { atomicAdd(address, val); }
+static inline __device__ void atomAdd(int64_t *address, int64_t val) {
+    AtomicAddIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
+}
+#if defined(USE_ROCM) || (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700 || CUDA_VERSION < 10000))
+static inline __device__ void atomAdd(at::Half *address, at::Half val) {
+    AtomicAddDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
+}
+#else
+static inline __device__ void atomAdd(at::Half *address, at::Half val) {
+    atomicAdd(reinterpret_cast<__half *>(address), val);
+}
+#endif
+static inline __device__ void atomAdd(float *address, float val) { atomicAdd(address, val); }
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600 || CUDA_VERSION < 8000)
+static inline __device__ void atomAdd(double *address, double val) {
+    AtomicAddDecimalImpl<double, sizeof(double)>()(address, val);
+}
+#else
+static inline __device__ void atomAdd(double *address, double val) { atomicAdd(address, val); }
+#endif
+static inline __device__ void atomAdd(at::BFloat16 *address, at::BFloat16 val) {
+    AtomicAddDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
+}
+
+#define OP(X, Y) Y *X
+ATOMIC(Mul)
+#undef OP
+static inline __device__ void atomMul(uint8_t *address, uint8_t val) {
+    AtomicMulIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
+}
+static inline __device__ void atomMul(int8_t *address, int8_t val) {
+    AtomicMulIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
+}
+static inline __device__ void atomMul(int16_t *address, int16_t val) {
+    AtomicMulIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
+}
+static inline __device__ void atomMul(int32_t *address, int32_t val) {
+    AtomicMulIntegerImpl<int32_t, sizeof(int32_t)>()(address, val);
+}
+static inline __device__ void atomMul(int64_t *address, int64_t val) {
+    AtomicMulIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
+}
+static inline __device__ void atomMul(float *address, float val) {
+    AtomicMulDecimalImpl<float, sizeof(float)>()(address, val);
+}
+static inline __device__ void atomMul(at::Half *address, at::Half val) {
+    AtomicMulDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
+}
+static inline __device__ void atomMul(double *address, double val) {
+    AtomicMulDecimalImpl<double, sizeof(double)>()(address, val);
+}
+static inline __device__ void atomMul(at::BFloat16 *address, at::BFloat16 val) {
+    AtomicMulDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
+}
+
+#define OP(X, Y) Y / X
+ATOMIC(Div)
+#undef OP
+static inline __device__ void atomDiv(uint8_t *address, uint8_t val) {
+    AtomicDivIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
+}
+static inline __device__ void atomDiv(int8_t *address, int8_t val) {
+    AtomicDivIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
+}
+static inline __device__ void atomDiv(int16_t *address, int16_t val) {
+    AtomicDivIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
+}
+static inline __device__ void atomDiv(int32_t *address, int32_t val) {
+    AtomicDivIntegerImpl<int32_t, sizeof(int32_t)>()(address, val);
+}
+static inline __device__ void atomDiv(int64_t *address, int64_t val) {
+    AtomicDivIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
+}
+static inline __device__ void atomDiv(at::Half *address, at::Half val) {
+    AtomicDivDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
+}
+static inline __device__ void atomDiv(float *address, float val) {
+    AtomicDivDecimalImpl<float, sizeof(float)>()(address, val);
+}
+static inline __device__ void atomDiv(double *address, double val) {
+    AtomicDivDecimalImpl<double, sizeof(double)>()(address, val);
+}
+static inline __device__ void atomDiv(at::BFloat16 *address, at::BFloat16 val) {
+    AtomicDivDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
+}
+
+#define OP(X, Y) max(Y, X)
+ATOMIC(Max)
+#undef OP
+static inline __device__ void atomMax(uint8_t *address, uint8_t val) {
+    AtomicMaxIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
+}
+static inline __device__ void atomMax(int8_t *address, int8_t val) {
+    AtomicMaxIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
+}
+static inline __device__ void atomMax(int16_t *address, int16_t val) {
+    AtomicMaxIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
+}
+static inline __device__ void atomMax(int32_t *address, int32_t val) { atomicMax(address, val); }
+static inline __device__ void atomMax(int64_t *address, int64_t val) {
+    AtomicMaxIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
+}
+static inline __device__ void atomMax(at::Half *address, at::Half val) {
+    AtomicMaxDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
+}
+static inline __device__ void atomMax(float *address, float val) {
+    AtomicMaxDecimalImpl<float, sizeof(float)>()(address, val);
+}
+static inline __device__ void atomMax(double *address, double val) {
+    AtomicMaxDecimalImpl<double, sizeof(double)>()(address, val);
+}
+static inline __device__ void atomMax(at::BFloat16 *address, at::BFloat16 val) {
+    AtomicMaxDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
+}
+
+#define OP(X, Y) min(Y, X)
+ATOMIC(Min)
+#undef OP
+static inline __device__ void atomMin(uint8_t *address, uint8_t val) {
+    AtomicMinIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
+}
+static inline __device__ void atomMin(int8_t *address, int8_t val) {
+    AtomicMinIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
+}
+static inline __device__ void atomMin(int16_t *address, int16_t val) {
+    AtomicMinIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
+}
+static inline __device__ void atomMin(int32_t *address, int32_t val) { atomicMin(address, val); }
+static inline __device__ void atomMin(int64_t *address, int64_t val) {
+    AtomicMinIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
+}
+static inline __device__ void atomMin(at::Half *address, at::Half val) {
+    AtomicMinDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
+}
+static inline __device__ void atomMin(float *address, float val) {
+    AtomicMinDecimalImpl<float, sizeof(float)>()(address, val);
+}
+static inline __device__ void atomMin(double *address, double val) {
+    AtomicMinDecimalImpl<double, sizeof(double)>()(address, val);
+}
+static inline __device__ void atomMin(at::BFloat16 *address, at::BFloat16 val) {
+    AtomicMinDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
+}

third_party/vipe/csrc/scatter_ext/cuda/reducer.cuh

@@ -0,0 +1,115 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#pragma once
+
+#include <limits>
+#include <map>
+
+#include "atomics.cuh"
+
+enum ReductionType { SUM, MEAN, MUL, DIV, MIN, MAX };
+
+const std::map<std::string, ReductionType> reduce2REDUCE = {
+    {"sum", SUM}, {"mean", MEAN}, {"mul", MUL}, {"div", DIV}, {"min", MIN}, {"max", MAX},
+};
+
+#define AT_DISPATCH_REDUCTION_TYPES(reduce, ...)              \
+    [&] {                                                     \
+        switch (reduce2REDUCE.at(reduce)) {                   \
+            case SUM: {                                       \
+                static constexpr ReductionType REDUCE = SUM;  \
+                return __VA_ARGS__();                         \
+            }                                                 \
+            case MEAN: {                                      \
+                static constexpr ReductionType REDUCE = MEAN; \
+                return __VA_ARGS__();                         \
+            }                                                 \
+            case MUL: {                                       \
+                static constexpr ReductionType REDUCE = MUL;  \
+                return __VA_ARGS__();                         \
+            }                                                 \
+            case DIV: {                                       \
+                static constexpr ReductionType REDUCE = DIV;  \
+                return __VA_ARGS__();                         \
+            }                                                 \
+            case MIN: {                                       \
+                static constexpr ReductionType REDUCE = MIN;  \
+                return __VA_ARGS__();                         \
+            }                                                 \
+            case MAX: {                                       \
+                static constexpr ReductionType REDUCE = MAX;  \
+                return __VA_ARGS__();                         \
+            }                                                 \
+        }                                                     \
+    }()
+
+template <typename scalar_t, ReductionType REDUCE>
+struct Reducer {
+    static inline __host__ __device__ scalar_t init() {
+        if (REDUCE == MUL || REDUCE == DIV)
+            return (scalar_t)1;
+        else if (REDUCE == MIN)
+            return std::numeric_limits<scalar_t>::max();
+        else if (REDUCE == MAX)
+            return std::numeric_limits<scalar_t>::lowest();
+        else
+            return (scalar_t)0;
+    }
+
+    static inline __host__ __device__ void update(scalar_t *val, scalar_t new_val) {
+        if (REDUCE == SUM || REDUCE == MEAN)
+            *val = *val + new_val;
+        else if (REDUCE == MUL)
+            *val = *val * new_val;
+        else if (REDUCE == DIV)
+            *val = *val / new_val;
+        else if ((REDUCE == MIN && new_val < *val) || (REDUCE == MAX && new_val > *val)) {
+            *val = new_val;
+        }
+    }
+
+    static inline __host__ __device__ void update(scalar_t *val, scalar_t new_val, int64_t *arg, int64_t new_arg) {
+        if (REDUCE == SUM || REDUCE == MEAN)
+            *val = *val + new_val;
+        else if (REDUCE == MUL)
+            *val = *val * new_val;
+        else if (REDUCE == DIV)
+            *val = *val / new_val;
+        else if ((REDUCE == MIN && new_val < *val) || (REDUCE == MAX && new_val > *val)) {
+            *val = new_val;
+            *arg = new_arg;
+        }
+    }
+
+    static inline __host__ __device__ void write(scalar_t *address, scalar_t val, int64_t *arg_address, int64_t arg,
+                                                 int count) {
+        if (REDUCE == SUM || REDUCE == MUL || REDUCE == DIV)
+            *address = val;
+        else if (REDUCE == MEAN)
+            *address = val / (scalar_t)(count > 0 ? count : 1);
+        else if (REDUCE == MIN || REDUCE == MAX) {
+            if (count > 0) {
+                *address = val;
+                *arg_address = arg;
+            } else
+                *address = (scalar_t)0;
+        }
+    }
+
+    static inline __device__ void atomic_write(scalar_t *address, scalar_t val) {
+        if (REDUCE == SUM || REDUCE == MEAN)
+            atomAdd(address, val);
+        else if (REDUCE == MUL)
+            atomMul(address, val);
+        else if (REDUCE == DIV)
+            atomDiv(address, val);
+        else if (REDUCE == MIN)
+            atomMin(address, val);
+        else if (REDUCE == MAX)
+            atomMax(address, val);
+    }
+};

third_party/vipe/csrc/scatter_ext/cuda/scatter_cuda.cu

@@ -0,0 +1,128 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include "scatter_cuda.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/Dispatch.h>
+#include <ATen/cuda/detail/IndexUtils.cuh>
+#include <ATen/cuda/detail/TensorInfo.cuh>
+
+#include "reducer.cuh"
+#include "utils.cuh"
+
+#define THREADS 256
+#define BLOCKS(N) (N + THREADS - 1) / THREADS
+
+template <typename scalar_t, ReductionType REDUCE>
+__global__ void scatter_kernel(const scalar_t *src_data, const at::cuda::detail::TensorInfo<int64_t, int> index_info,
+                               scalar_t *out_data, int E, int K, int N, int numel) {
+    int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
+
+    int b = thread_idx / (E * K);
+    int k = thread_idx % K;
+
+    if (thread_idx < numel) {
+        int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(thread_idx, index_info);
+        int64_t idx = index_info.data[offset];
+
+        Reducer<scalar_t, REDUCE>::atomic_write(out_data + b * N * K + idx * K + k, src_data[thread_idx]);
+    }
+}
+
+template <typename scalar_t>
+__global__ void scatter_arg_kernel(const scalar_t *src_data,
+                                   const at::cuda::detail::TensorInfo<int64_t, int> index_info,
+                                   const scalar_t *out_data, int64_t *arg_out_data, int E, int K, int N, int numel) {
+    int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
+
+    int b = thread_idx / (E * K);
+    int e = (thread_idx / K) % E;
+    int k = thread_idx % K;
+
+    if (thread_idx < numel) {
+        int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(thread_idx, index_info);
+        int64_t idx = index_info.data[offset];
+
+        if (src_data[thread_idx] == out_data[b * N * K + idx * K + k]) {
+            arg_out_data[b * N * K + idx * K + k] = e;
+        }
+    }
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_cuda(torch::Tensor src, torch::Tensor index,
+                                                                     int64_t dim,
+                                                                     std::optional<torch::Tensor> optional_out,
+                                                                     std::optional<int64_t> dim_size,
+                                                                     std::string reduce) {
+    CHECK_CUDA(src);
+    CHECK_CUDA(index);
+    if (optional_out.has_value()) CHECK_CUDA(optional_out.value());
+    c10::cuda::MaybeSetDevice(src.get_device());
+
+    CHECK_INPUT(src.dim() == index.dim());
+    for (auto i = 0; i < index.dim() - 1; i++) CHECK_INPUT(src.size(i) >= index.size(i));
+
+    src = src.contiguous();
+
+    torch::Tensor out;
+    if (optional_out.has_value()) {
+        out = optional_out.value().contiguous();
+        for (auto i = 0; i < out.dim(); i++)
+            if (i != dim) CHECK_INPUT(src.size(i) == out.size(i));
+    } else {
+        auto sizes = src.sizes().vec();
+        if (dim_size.has_value())
+            sizes[dim] = dim_size.value();
+        else if (index.numel() == 0)
+            sizes[dim] = 0;
+        else {
+            sizes[dim] = 1 + index.max().cpu().data_ptr<int64_t>()[0];
+        }
+        out = torch::empty(sizes, src.options());
+    }
+
+    std::optional<torch::Tensor> arg_out = std::nullopt;
+    int64_t *arg_out_data = nullptr;
+    if (reduce2REDUCE.at(reduce) == MIN || reduce2REDUCE.at(reduce) == MAX) {
+        arg_out = torch::full_like(out, src.size(dim), index.options());
+        arg_out_data = arg_out.value().data_ptr<int64_t>();
+    }
+
+    if (src.numel() == 0) {
+        if (!optional_out.has_value()) out.fill_(0);
+        return std::make_tuple(out, arg_out);
+    }
+
+    auto B = 1;
+    for (auto i = 0; i < dim; i++) B *= src.size(i);
+    auto E = src.size(dim);
+    auto K = src.numel() / (B * E);
+    auto N = out.size(dim);
+
+    auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
+    auto stream = at::cuda::getCurrentCUDAStream();
+    AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
+        auto src_data = src.data_ptr<scalar_t>();
+        auto out_data = out.data_ptr<scalar_t>();
+
+        AT_DISPATCH_REDUCTION_TYPES(reduce, [&] {
+            if (!optional_out.has_value()) out.fill_(Reducer<scalar_t, REDUCE>::init());
+
+            scatter_kernel<scalar_t, REDUCE>
+                <<<BLOCKS(src.numel()), THREADS, 0, stream>>>(src_data, index_info, out_data, E, K, N, src.numel());
+
+            if (!optional_out.has_value() && (REDUCE == MIN || REDUCE == MAX))
+                out.masked_fill_(out == Reducer<scalar_t, REDUCE>::init(), (scalar_t)0);
+
+            if (REDUCE == MIN || REDUCE == MAX)
+                scatter_arg_kernel<scalar_t><<<BLOCKS(src.numel()), THREADS, 0, stream>>>(
+                    src_data, index_info, out_data, arg_out_data, E, K, N, src.numel());
+        });
+    });
+
+    return std::make_tuple(out, arg_out);
+}

third_party/vipe/csrc/scatter_ext/cuda/scatter_cuda.h

@@ -0,0 +1,15 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#pragma once
+
+#include <torch/extension.h>
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_cuda(torch::Tensor src, torch::Tensor index,
+                                                                     int64_t dim,
+                                                                     std::optional<torch::Tensor> optional_out,
+                                                                     std::optional<int64_t> dim_size,
+                                                                     std::string reduce);

third_party/vipe/csrc/scatter_ext/cuda/utils.cuh

@@ -0,0 +1,29 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#pragma once
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor")
+#define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch")
+
+__device__ __inline__ at::Half __shfl_up_sync(const unsigned mask, const at::Half var, const unsigned int delta) {
+    return __shfl_up_sync(mask, var.operator __half(), delta);
+}
+
+__device__ __inline__ at::Half __shfl_down_sync(const unsigned mask, const at::Half var, const unsigned int delta) {
+    return __shfl_down_sync(mask, var.operator __half(), delta);
+}
+
+__device__ __inline__ at::Half __shfl_up(const at::Half var, const unsigned int delta) {
+    return __shfl_up(var.operator __half(), delta);
+}
+
+__device__ __inline__ at::Half __shfl_down(const at::Half var, const unsigned int delta) {
+    return __shfl_down(var.operator __half(), delta);
+}
+
+#define SHFL_UP_SYNC __shfl_up_sync
+#define SHFL_DOWN_SYNC __shfl_down_sync

third_party/vipe/csrc/scatter_ext/scatter.cpp

@@ -0,0 +1,238 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <vector>
+
+#include "cuda/scatter_cuda.h"
+
+inline std::vector<int64_t> list2vec(const c10::List<int64_t> list) {
+    std::vector<int64_t> result;
+    result.reserve(list.size());
+    for (size_t i = 0; i < list.size(); i++) result.push_back(list[i]);
+    return result;
+}
+
+torch::Tensor broadcast(torch::Tensor src, torch::Tensor other, int64_t dim) {
+    if (src.dim() == 1)
+        for (auto i = 0; i < dim; i++) src = src.unsqueeze(0);
+    for (auto i = src.dim(); i < other.dim(); i++) src = src.unsqueeze(-1);
+    src = src.expand(other.sizes().vec());
+    return src;
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_fw(torch::Tensor src, torch::Tensor index, int64_t dim,
+                                                                   std::optional<torch::Tensor> optional_out,
+                                                                   std::optional<int64_t> dim_size,
+                                                                   std::string reduce) {
+    return scatter_cuda(src, index, dim, optional_out, dim_size, reduce);
+}
+
+using torch::autograd::AutogradContext;
+using torch::autograd::Variable;
+using torch::autograd::variable_list;
+
+class ScatterSum : public torch::autograd::Function<ScatterSum> {
+   public:
+    static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
+                                 std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
+        dim = dim < 0 ? src.dim() + dim : dim;
+        ctx->saved_data["dim"] = dim;
+        ctx->saved_data["src_shape"] = src.sizes();
+        index = broadcast(index, src, dim);
+        auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");
+        auto out = std::get<0>(result);
+        ctx->save_for_backward({index});
+        if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
+        return {out};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
+        auto grad_out = grad_outs[0];
+        auto saved = ctx->get_saved_variables();
+        auto index = saved[0];
+        auto dim = ctx->saved_data["dim"].toInt();
+        auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
+        auto grad_in = torch::gather(grad_out, dim, index, false);
+        return {grad_in, Variable(), Variable(), Variable(), Variable()};
+    }
+};
+
+class ScatterMul : public torch::autograd::Function<ScatterMul> {
+   public:
+    static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
+                                 std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
+        dim = dim < 0 ? src.dim() + dim : dim;
+        ctx->saved_data["dim"] = dim;
+        ctx->saved_data["src_shape"] = src.sizes();
+        index = broadcast(index, src, dim);
+        auto result = scatter_fw(src, index, dim, optional_out, dim_size, "mul");
+        auto out = std::get<0>(result);
+        ctx->save_for_backward({src, index, out});
+        if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
+        return {out};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
+        auto grad_out = grad_outs[0];
+        auto saved = ctx->get_saved_variables();
+        auto src = saved[0];
+        auto index = saved[1];
+        auto out = saved[2];
+        auto dim = ctx->saved_data["dim"].toInt();
+        auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
+        auto grad_in = torch::gather(grad_out * out, dim, index, false).div_(src);
+        grad_in.masked_fill_(grad_in.isnan(), 0);
+        return {grad_in, Variable(), Variable(), Variable(), Variable()};
+    }
+};
+
+class ScatterMean : public torch::autograd::Function<ScatterMean> {
+   public:
+    static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
+                                 std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
+        dim = dim < 0 ? src.dim() + dim : dim;
+        ctx->saved_data["dim"] = dim;
+        ctx->saved_data["src_shape"] = src.sizes();
+
+        auto old_index = index;
+
+        index = broadcast(index, src, dim);
+        auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");
+        auto out = std::get<0>(result);
+
+        auto ones = torch::ones(old_index.sizes(), src.options());
+        result = scatter_fw(ones, old_index, old_index.dim() <= dim ? old_index.dim() - 1 : dim, std::nullopt,
+                            out.size(dim), "sum");
+        auto count = std::get<0>(result);
+        count.masked_fill_(count < 1, 1);
+        count = broadcast(count, out, dim);
+        if (out.is_floating_point())
+            out.true_divide_(count);
+        else
+            out.div_(count, "floor");
+
+        ctx->save_for_backward({index, count});
+        if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
+        return {out};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
+        auto grad_out = grad_outs[0];
+        auto saved = ctx->get_saved_variables();
+        auto index = saved[0];
+        auto count = saved[1];
+        auto dim = ctx->saved_data["dim"].toInt();
+        auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
+        count = torch::gather(count, dim, index, false);
+        auto grad_in = torch::gather(grad_out, dim, index, false);
+        grad_in.true_divide_(count);
+        return {grad_in, Variable(), Variable(), Variable(), Variable()};
+    }
+};
+
+class ScatterMin : public torch::autograd::Function<ScatterMin> {
+   public:
+    static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
+                                 std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
+        dim = dim < 0 ? src.dim() + dim : dim;
+        ctx->saved_data["dim"] = dim;
+        ctx->saved_data["src_shape"] = src.sizes();
+
+        index = broadcast(index, src, dim);
+        auto result = scatter_fw(src, index, dim, optional_out, dim_size, "min");
+        auto out = std::get<0>(result);
+        auto arg_out = std::get<1>(result).value();
+        ctx->save_for_backward({index, arg_out});
+        ctx->mark_non_differentiable({arg_out});
+        if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
+        return {out, arg_out};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
+        auto grad_out = grad_outs[0];
+        auto saved = ctx->get_saved_variables();
+        auto index = saved[0];
+        auto arg_out = saved[1];
+        auto dim = ctx->saved_data["dim"].toInt();
+        auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
+        src_shape[dim] += 1;
+        auto grad_in = torch::zeros(src_shape, grad_out.options());
+        grad_in.scatter_(dim, arg_out, grad_out);
+        grad_in = grad_in.narrow(dim, 0, src_shape[dim] - 1);
+        return {grad_in, Variable(), Variable(), Variable(), Variable()};
+    }
+};
+
+class ScatterMax : public torch::autograd::Function<ScatterMax> {
+   public:
+    static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
+                                 std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
+        dim = dim < 0 ? src.dim() + dim : dim;
+        ctx->saved_data["dim"] = dim;
+        ctx->saved_data["src_shape"] = src.sizes();
+
+        index = broadcast(index, src, dim);
+        auto result = scatter_fw(src, index, dim, optional_out, dim_size, "max");
+        auto out = std::get<0>(result);
+        auto arg_out = std::get<1>(result).value();
+        ctx->save_for_backward({index, arg_out});
+        ctx->mark_non_differentiable({arg_out});
+        if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
+        return {out, arg_out};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
+        auto grad_out = grad_outs[0];
+        auto saved = ctx->get_saved_variables();
+        auto index = saved[0];
+        auto arg_out = saved[1];
+        auto dim = ctx->saved_data["dim"].toInt();
+        auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
+        src_shape[dim] += 1;
+        auto grad_in = torch::zeros(src_shape, grad_out.options());
+        grad_in.scatter_(dim, arg_out, grad_out);
+        grad_in = grad_in.narrow(dim, 0, src_shape[dim] - 1);
+        return {grad_in, Variable(), Variable(), Variable(), Variable()};
+    }
+};
+
+torch::Tensor scatter_sum(torch::Tensor src, torch::Tensor index, int64_t dim,
+                          std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
+    return ScatterSum::apply(src, index, dim, optional_out, dim_size)[0];
+}
+
+torch::Tensor scatter_mul(torch::Tensor src, torch::Tensor index, int64_t dim,
+                          std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
+    return ScatterMul::apply(src, index, dim, optional_out, dim_size)[0];
+}
+
+torch::Tensor scatter_mean(torch::Tensor src, torch::Tensor index, int64_t dim,
+                           std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
+    return ScatterMean::apply(src, index, dim, optional_out, dim_size)[0];
+}
+
+std::tuple<torch::Tensor, torch::Tensor> scatter_min(torch::Tensor src, torch::Tensor index, int64_t dim,
+                                                     std::optional<torch::Tensor> optional_out,
+                                                     std::optional<int64_t> dim_size) {
+    auto result = ScatterMin::apply(src, index, dim, optional_out, dim_size);
+    return std::make_tuple(result[0], result[1]);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> scatter_max(torch::Tensor src, torch::Tensor index, int64_t dim,
+                                                     std::optional<torch::Tensor> optional_out,
+                                                     std::optional<int64_t> dim_size) {
+    auto result = ScatterMax::apply(src, index, dim, optional_out, dim_size);
+    return std::make_tuple(result[0], result[1]);
+}
+
+void pybind_scatter_ext(py::module &m) {
+    m.def("scatter_sum", &scatter_sum, "Scatter sum operation");
+    m.def("scatter_mul", &scatter_mul, "Scatter multiplication operation");
+    m.def("scatter_mean", &scatter_mean, "Scatter mean operation");
+    m.def("scatter_min", &scatter_min, "Scatter min operation");
+    m.def("scatter_max", &scatter_max, "Scatter max operation");
+}

third_party/vipe/csrc/slam_ext/slam.cpp

@@ -0,0 +1,37 @@
+/**
+ * This file includes code originally from the DROID-SLAM repository:
+ * https://github.com/princeton-vl/DROID-SLAM
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <vector>
+
+namespace slam_ext {
+
+torch::Tensor depth_filter_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics, torch::Tensor ix,
+                                torch::Tensor thresh);
+
+torch::Tensor frame_distance_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics, torch::Tensor pi,
+                                  torch::Tensor pj, torch::Tensor qi, torch::Tensor qj, torch::Tensor di,
+                                  const float beta);
+
+std::vector<torch::Tensor> projmap_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics,
+                                        torch::Tensor ii, torch::Tensor jj);
+
+torch::Tensor iproj_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics);
+
+std::vector<torch::Tensor> ba_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics,
+                                   torch::Tensor disps_sens, torch::Tensor targets, torch::Tensor weights,
+                                   torch::Tensor eta, torch::Tensor ii, torch::Tensor jj, const int t0, const int t1,
+                                   const int iterations, const float lm, const float ep, const bool motion_only);
+
+}  // namespace slam_ext
+
+void pybind_slam_ext(py::module &m) {
+    m.def("ba", &slam_ext::ba_cuda, "bundle adjustment");
+    m.def("frame_distance", &slam_ext::frame_distance_cuda, "frame_distance");
+    m.def("projmap", &slam_ext::projmap_cuda, "projmap");
+    m.def("depth_filter", &slam_ext::depth_filter_cuda, "depth_filter");
+    m.def("iproj", &slam_ext::iproj_cuda, "back projection");
+}

third_party/vipe/csrc/utils_ext/cuda_kdtree.cuh

@@ -0,0 +1,137 @@
+/***********************************************************************
+ * The code in this file is modified from the original code in the
+ * https://github.com/jlblancoc/nanoflann repository.
+ *
+ * Software License Agreement (BSD License)
+ *
+ * Copyright 2008-2009  Marius Muja (mariusm@cs.ubc.ca). All rights reserved.
+ * Copyright 2008-2009  David G. Lowe (lowe@cs.ubc.ca). All rights reserved.
+ * Copyright 2011-2025  Jose Luis Blanco (joseluisblancoc@gmail.com).
+ *   All rights reserved.
+ *
+ * THE BSD LICENSE
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
+ * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
+ * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
+ * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
+ * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *************************************************************************/
+
+#ifndef FLANN_CUDA_KD_TREE_BUILDER_H_
+#define FLANN_CUDA_KD_TREE_BUILDER_H_
+
+#include <thrust/host_vector.h>
+#include <thrust/device_vector.h>
+#include <thrust/sort.h>
+#include <thrust/partition.h>
+#include <thrust/unique.h>
+#include <thrust/scan.h>
+#include "math_util.h"
+#include <cstdlib>
+#include <map>
+
+namespace tinyflann {
+
+// Distance types.
+struct CudaL1;
+struct CudaL2;
+
+// Parameters
+struct KDTreeCuda3dIndexParams {
+    int leaf_max_size = 64;
+};
+
+struct SearchParams {
+    SearchParams(int checks_ = 32, float eps_ = 0.0, bool sorted_ = true)
+        : checks(checks_), eps(eps_), sorted(sorted_) {
+        max_neighbors = -1;
+        use_heap = true;
+    }
+
+    // how many leafs to visit when searching for neighbours (-1 for unlimited)
+    int checks;
+    // search for eps-approximate neighbours (default: 0)
+    float eps;
+    // only for radius search, require neighbours sorted by distance (default: true)
+    bool sorted;
+    // maximum number of neighbors radius search should return (-1 for unlimited)
+    int max_neighbors;
+    // use a heap to manage the result set (default: FLANN_Undefined)
+    bool use_heap;
+};
+
+template <typename Distance>
+class KDTreeCuda3dIndex {
+   public:
+    int visited_leafs;
+    KDTreeCuda3dIndex(const float* input_data, size_t input_count,
+                      const KDTreeCuda3dIndexParams& params = KDTreeCuda3dIndexParams())
+        : dataset_(input_data), leaf_count_(0), visited_leafs(0), node_count_(0), current_node_count_(0) {
+        size_ = input_count;
+        leaf_max_size_ = params.leaf_max_size;
+        gpu_helper_ = 0;
+    }
+
+    /**
+     * Standard destructor
+     */
+    ~KDTreeCuda3dIndex() { clearGpuBuffers(); }
+
+    /**
+     * Builds the index
+     */
+    void buildIndex() {
+        leaf_count_ = 0;
+        node_count_ = 0;
+        uploadTreeToGpu();
+    }
+
+    /**
+     * queries: (N, p) flattened float cuda array where only first 3 elements are used.
+     * n_query: N
+     * n_query_stride: p
+     * indices: (N, knn) int cuda array
+     * dists: (N, knn) float cuda array
+     */
+    void knnSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, size_t knn,
+                   const SearchParams& params = SearchParams()) const;
+    int radiusSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, float radius,
+                     const SearchParams& params = SearchParams()) const;
+
+   private:
+    void uploadTreeToGpu();
+    void clearGpuBuffers();
+
+   private:
+    struct GpuHelper;
+    GpuHelper* gpu_helper_;
+
+    const float* dataset_;
+    int leaf_max_size_;
+    int leaf_count_;
+    int node_count_;
+    //! used by convertTreeToGpuFormat
+    int current_node_count_;
+    size_t size_;
+
+};  // class KDTreeCuda3dIndex
+
+}  // namespace tinyflann
+
+#endif

third_party/vipe/csrc/utils_ext/knn.cu

@@ -0,0 +1,66 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAException.h>
+
+#include "cuda_kdtree.cuh"
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_IS_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, #x " must be a float tensor")
+
+std::vector<torch::Tensor> nearestNeighbours(torch::Tensor query, torch::Tensor tree, int knn) {
+    CHECK_CUDA(query);
+    CHECK_IS_FLOAT(query);
+    CHECK_CUDA(tree);
+    CHECK_IS_FLOAT(tree);
+    TORCH_CHECK(tree.size(0) >= knn, "knn is too small compared to the size of point cloud!");
+
+    torch::Tensor strided_tree = tree;
+    torch::Tensor strided_query = query;
+    torch::Device device = tree.device();
+    long n_queries = query.size(0);
+    long n_tree = tree.size(0);
+
+    if (strided_tree.stride(0) != 4) {
+        strided_tree = torch::zeros({n_tree, 4}, torch::dtype(torch::kFloat32).device(device));
+        strided_tree.index_put_(
+            {torch::indexing::Ellipsis, torch::indexing::Slice(torch::indexing::None, tree.size(1))}, tree);
+    }
+
+    if (strided_query.stride(0) != 3) {
+        strided_query = torch::zeros({n_queries, 3}, torch::dtype(torch::kFloat32).device(device));
+        strided_query.index_put_(
+            {torch::indexing::Ellipsis, torch::indexing::Slice(torch::indexing::None, query.size(1))}, query);
+    }
+
+    auto *knn_index = new tinyflann::KDTreeCuda3dIndex<tinyflann::CudaL2>(strided_tree.data_ptr<float>(), n_tree);
+    knn_index->buildIndex();
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    torch::Tensor dist = torch::empty(n_queries * knn, torch::dtype(torch::kFloat32).device(device));
+    torch::Tensor indices = torch::empty(n_queries * knn, torch::dtype(torch::kInt32).device(device));
+
+    knn_index->knnSearch(strided_query.data_ptr<float>(), n_queries, 3, indices.data_ptr<int>(), dist.data_ptr<float>(),
+                         knn);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    delete knn_index;
+
+    return {dist.reshape({n_queries, knn}), indices.reshape({n_queries, knn})};
+}

third_party/vipe/csrc/utils_ext/platform.h

@@ -0,0 +1,50 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
+#define _CPU_AND_GPU_CODE_ __host__ __device__  // for CUDA device code
+#else
+#define _CPU_AND_GPU_CODE_
+#endif
+
+#if defined(__CUDACC__)
+#define _CPU_AND_GPU_CODE_TEMPLATE_ __host__ __device__  // for CUDA device code
+#else
+#define _CPU_AND_GPU_CODE_TEMPLATE_
+#endif
+
+#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
+#define _CPU_AND_GPU_CONSTANT_ __constant__  // for CUDA device code
+#else
+#define _CPU_AND_GPU_CONSTANT_ const
+#endif
+
+#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
+#include <THC/THCAtomics.cuh>
+
+template <typename ScalarT>
+static inline __device__ void atomicAdd(ScalarT* addr, ScalarT value) {
+    gpuAtomicAddNoReturn(addr, value);
+}
+#else
+template <typename ScalarT>
+static inline void atomicAdd(ScalarT* addr, ScalarT value) {
+    *addr += value;
+}
+#endif

third_party/vipe/csrc/utils_ext/utils_bind.cpp

@@ -0,0 +1,25 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/extension.h>
+
+std::vector<torch::Tensor> nearestNeighbours(torch::Tensor query, torch::Tensor tree, int knn);
+
+void pybind_utils_ext(py::module &m) {
+    m.def("nearest_neighbours", &nearestNeighbours, "KNN computation (CUDA) distance is squared L2.", py::arg("query"),
+          py::arg("tree"), py::arg("knn"));
+}

third_party/vipe/envs/base.yml

@@ -0,0 +1,14 @@
+name: vipe
+channels:
+  - nvidia/label/cuda-12.8.0
+  - conda-forge
+  - nodefaults
+dependencies:
+  - python=3.10
+  - pip
+  - cuda-nvcc
+  - eigen
+  - zlib
+  - libcusparse-dev
+  - libcublas-dev
+  - libcusolver-dev

third_party/vipe/envs/requirements.in

@@ -0,0 +1,22 @@
+torch==2.7.0+cu128
+torchvision
+
+click
+numpy
+einops
+omegaconf
+tqdm
+rerun-sdk
+imageio[ffmpeg]
+opencv-python
+matplotlib
+python-pycg
+timm
+transformers
+gdown
+OpenEXR<3.3.0
+kornia
+hydra-core
+ray
+ninja
+viser

third_party/vipe/envs/requirements.txt

@@ -0,0 +1,373 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile --emit-index-url --prerelease=allow -o envs/requirements.txt envs/requirements.in
+
+aiosignal==1.3.1
+    # via ray
+antlr4-python3-runtime==4.9.3
+    # via
+    #   hydra-core
+    #   omegaconf
+attrs==24.2.0
+    # via
+    #   jsonschema
+    #   referencing
+    #   rerun-sdk
+beautifulsoup4==4.13.3
+    # via gdown
+calmsize==0.1.3
+    # via python-pycg
+certifi==2022.12.7
+    # via requests
+charset-normalizer==2.1.1
+    # via
+    #   requests
+    #   trimesh
+click==8.1.8
+    # via ray
+colorlog==6.9.0
+    # via trimesh
+contourpy==1.2.1
+    # via matplotlib
+cycler==0.12.1
+    # via matplotlib
+einops==0.8.1
+    # via -r envs/requirements.in
+embreex==2.17.7.post6
+    # via trimesh
+filelock==3.13.1
+    # via
+    #   gdown
+    #   huggingface-hub
+    #   ray
+    #   torch
+    #   transformers
+fonttools==4.53.0
+    # via matplotlib
+frozenlist==1.4.1
+    # via
+    #   aiosignal
+    #   ray
+fsspec==2024.6.1
+    # via
+    #   huggingface-hub
+    #   torch
+gdown==5.2.0
+    # via -r envs/requirements.in
+h11==0.16.0
+    # via httpx
+httpx==1.0.dev1
+    # via trimesh
+huggingface-hub==0.25.1
+    # via
+    #   timm
+    #   tokenizers
+    #   transformers
+hydra-core==1.3.2
+    # via -r envs/requirements.in
+idna==3.4
+    # via requests
+imageio==2.37.0
+    # via
+    #   -r envs/requirements.in
+    #   python-pycg
+    #   scikit-image
+    #   viser
+imageio-ffmpeg==0.6.0
+    # via imageio
+jinja2==3.1.4
+    # via torch
+jsonschema==4.23.0
+    # via
+    #   ray
+    #   trimesh
+jsonschema-specifications==2024.10.1
+    # via jsonschema
+kiwisolver==1.4.5
+    # via matplotlib
+kornia==0.8.0
+    # via -r envs/requirements.in
+kornia-rs==0.1.8
+    # via kornia
+lazy-loader==0.4
+    # via scikit-image
+lxml==6.0.0
+    # via
+    #   trimesh
+    #   yourdfpy
+manifold3d==3.2.0
+    # via trimesh
+mapbox-earcut==1.0.3
+    # via trimesh
+markdown-it-py==3.0.0
+    # via rich
+markupsafe==2.1.5
+    # via jinja2
+matplotlib==3.9.0
+    # via
+    #   -r envs/requirements.in
+    #   python-pycg
+mdurl==0.1.2
+    # via markdown-it-py
+mpmath==1.3.0
+    # via sympy
+msgpack==1.1.0
+    # via ray
+msgspec==0.19.0
+    # via viser
+networkx==3.3
+    # via
+    #   scikit-image
+    #   torch
+    #   trimesh
+ninja==1.11.1.4
+    # via -r envs/requirements.in
+nodeenv==1.9.1
+    # via viser
+numpy==2.1.2
+    # via
+    #   -r envs/requirements.in
+    #   contourpy
+    #   embreex
+    #   imageio
+    #   manifold3d
+    #   mapbox-earcut
+    #   matplotlib
+    #   opencv-python
+    #   plyfile
+    #   pyarrow
+    #   pycollada
+    #   pyquaternion
+    #   python-pycg
+    #   rerun-sdk
+    #   scikit-image
+    #   scipy
+    #   shapely
+    #   tifffile
+    #   torchvision
+    #   transformers
+    #   trimesh
+    #   vhacdx
+    #   viser
+    #   yourdfpy
+nvidia-cublas-cu12==12.8.3.14
+    # via
+    #   nvidia-cudnn-cu12
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-cuda-cupti-cu12==12.8.57
+    # via torch
+nvidia-cuda-nvrtc-cu12==12.8.61
+    # via torch
+nvidia-cuda-runtime-cu12==12.8.57
+    # via torch
+nvidia-cudnn-cu12==9.7.1.26
+    # via torch
+nvidia-cufft-cu12==11.3.3.41
+    # via torch
+nvidia-cufile-cu12==1.13.0.11
+    # via torch
+nvidia-curand-cu12==10.3.9.55
+    # via torch
+nvidia-cusolver-cu12==11.7.2.55
+    # via torch
+nvidia-cusparse-cu12==12.5.7.53
+    # via
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-cusparselt-cu12==0.6.3
+    # via torch
+nvidia-ml-py==12.570.86
+    # via pynvml
+nvidia-nccl-cu12==2.26.2
+    # via torch
+nvidia-nvjitlink-cu12==12.8.61
+    # via
+    #   nvidia-cufft-cu12
+    #   nvidia-cusolver-cu12
+    #   nvidia-cusparse-cu12
+    #   torch
+nvidia-nvtx-cu12==12.8.55
+    # via torch
+omegaconf==2.3.0
+    # via
+    #   -r envs/requirements.in
+    #   hydra-core
+    #   python-pycg
+opencv-python==4.11.0.86
+    # via
+    #   -r envs/requirements.in
+    #   viser
+openexr==3.2.4
+    # via -r envs/requirements.in
+packaging==24.1
+    # via
+    #   huggingface-hub
+    #   hydra-core
+    #   kornia
+    #   lazy-loader
+    #   matplotlib
+    #   ray
+    #   scikit-image
+    #   transformers
+pillow==11.0.0
+    # via
+    #   imageio
+    #   matplotlib
+    #   python-pycg
+    #   rerun-sdk
+    #   scikit-image
+    #   torchvision
+    #   trimesh
+plyfile==1.1
+    # via python-pycg
+protobuf==6.30.0rc1
+    # via ray
+psutil==6.0.0
+    # via
+    #   imageio
+    #   viser
+pyarrow==17.0.0
+    # via rerun-sdk
+pycollada==0.9.2
+    # via trimesh
+pygments==2.19.1
+    # via rich
+pynvml==12.0.0
+    # via python-pycg
+pyparsing==3.1.2
+    # via matplotlib
+pyquaternion==0.9.9
+    # via python-pycg
+pysocks==1.7.1
+    # via requests
+python-dateutil==2.9.0
+    # via
+    #   matplotlib
+    #   pycollada
+python-pycg==0.10.1
+    # via -r envs/requirements.in
+pyyaml==6.0.2
+    # via
+    #   huggingface-hub
+    #   omegaconf
+    #   python-pycg
+    #   ray
+    #   timm
+    #   transformers
+ray==2.42.1
+    # via -r envs/requirements.in
+referencing==0.36.2
+    # via
+    #   jsonschema
+    #   jsonschema-specifications
+regex==2024.9.11
+    # via transformers
+requests==2.28.1
+    # via
+    #   gdown
+    #   huggingface-hub
+    #   ray
+    #   transformers
+    #   viser
+rerun-sdk==0.22.0
+    # via -r envs/requirements.in
+rich==13.9.4
+    # via
+    #   python-pycg
+    #   viser
+rpds-py==0.22.3
+    # via
+    #   jsonschema
+    #   referencing
+rtree==1.4.0
+    # via trimesh
+safetensors==0.4.5
+    # via
+    #   timm
+    #   transformers
+scikit-image==0.25.2
+    # via viser
+scipy==1.15.1
+    # via
+    #   python-pycg
+    #   scikit-image
+    #   trimesh
+    #   viser
+screeninfo==0.8.1
+    # via python-pycg
+setuptools==70.2.0
+    # via
+    #   calmsize
+    #   trimesh
+    #   triton
+shapely==2.1.1
+    # via trimesh
+six==1.16.0
+    # via
+    #   python-dateutil
+    #   yourdfpy
+soupsieve==2.6
+    # via beautifulsoup4
+svg-path==7.0
+    # via trimesh
+sympy==1.13.3
+    # via torch
+tifffile==2025.5.10
+    # via scikit-image
+timm==1.0.14
+    # via -r envs/requirements.in
+tokenizers==0.21.0
+    # via transformers
+torch==2.7.0+cu128
+    # via
+    #   -r envs/requirements.in
+    #   kornia
+    #   timm
+    #   torchvision
+torchvision==0.22.0+cu128
+    # via
+    #   -r envs/requirements.in
+    #   timm
+tqdm==4.66.5
+    # via
+    #   -r envs/requirements.in
+    #   gdown
+    #   huggingface-hub
+    #   python-pycg
+    #   transformers
+    #   viser
+transformers==4.48.3
+    # via -r envs/requirements.in
+trimesh==4.6.2
+    # via
+    #   python-pycg
+    #   viser
+    #   yourdfpy
+triton==3.3.0
+    # via torch
+truststore==0.10.0
+    # via httpx
+typing-extensions==4.12.2
+    # via
+    #   beautifulsoup4
+    #   huggingface-hub
+    #   referencing
+    #   rerun-sdk
+    #   rich
+    #   torch
+urllib3==1.26.13
+    # via requests
+usd-core==25.2.post1
+    # via python-pycg
+vhacdx==0.0.8.post2
+    # via trimesh
+viser==1.0.1
+    # via -r envs/requirements.in
+websockets==15.0.1
+    # via viser
+xatlas==0.0.11
+    # via trimesh
+xxhash==3.5.0
+    # via trimesh
+yourdfpy==0.0.58
+    # via viser

third_party/vipe/pyproject.toml

@@ -0,0 +1,46 @@
+[build-system]
+requires = ["setuptools>=40.8.0", "wheel", "torch"]
+build-backend = "setuptools.build_meta"
+
+[tool.ruff]
+target-version = "py311"
+line-length = 120
+
+[tool.ruff.lint.isort]
+lines-between-types = 1
+lines-after-imports = 2
+
+[project]
+name = "vipe"
+description = "NVIDIA Video Pose Engine"
+authors = [
+    {name = "The ViPE Authors", email = "jiahuih@nvidia.com"},
+]
+readme = "README.md"
+requires-python = ">=3.10"
+dynamic = ["version"]
+dependencies = [
+    "numpy",
+    "einops",
+    "omegaconf",
+    "tqdm",
+    "rerun-sdk",
+    "imageio[ffmpeg]",
+    "opencv-python",
+    "python-pycg",
+    "timm",
+    "transformers",
+    "gdown",
+    "OpenEXR",
+    "kornia",
+    "viser",
+    "click",
+]
+
+[project.optional-dependencies]
+dav3 = [
+    "depth-anything-3 @ git+https://github.com/heiwang1997/Depth-Anything-3.git@main"
+]
+
+[project.scripts]
+vipe = "vipe.cli.main:main"

third_party/vipe/run.py

@@ -0,0 +1,28 @@
+import hydra
+from omegaconf import DictConfig
+
+
+@hydra.main(version_base=None, config_path="configs", config_name="default")
+def run(args: DictConfig) -> None:
+    from vipe.streams.base import StreamList
+
+    # Gather all video streams
+    stream_list = StreamList.make(args.streams)
+
+    from vipe.pipeline import make_pipeline
+    from vipe.utils.logging import configure_logging
+
+    # Process each video stream
+    logger = configure_logging()
+    for stream_idx in range(len(stream_list)):
+        video_stream = stream_list[stream_idx]
+        logger.info(
+            f"Processing {video_stream.name()} ({stream_idx + 1} / {len(stream_list)})"
+        )
+        pipeline = make_pipeline(args.pipeline)
+        pipeline.run(video_stream)
+        logger.info(f"Finished processing {video_stream.name()}")
+
+
+if __name__ == "__main__":
+    run()

third_party/vipe/scripts/download_dataset.py

@@ -0,0 +1,164 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+import shutil
+import tarfile
+import tempfile
+import zipfile
+
+from pathlib import Path
+
+import gdown
+import pandas as pd
+
+from huggingface_hub import HfApi
+
+
+def download_clips_from_url(
+    url: str, clips_timestamps: list[tuple[str, str]], output_paths: list[Path], cookies_from_browser: bool = False
+):
+    import datetime
+
+    import ffmpeg
+    import yt_dlp
+
+    def _get_seconds(t: str) -> float:
+        time_format = "%H:%M:%S.%f"
+        t_obj = datetime.datetime.strptime(t, time_format).time()
+        return t_obj.second + t_obj.microsecond / 1e6 + t_obj.minute * 60 + t_obj.hour * 3600
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        video_path = Path(tmpdir) / "video.mp4"
+
+        ydl_opts = {
+            "outtmpl": str(video_path),
+            "format": "wv*[height>=720][ext=mp4]/w[height>=720][ext=mp4]/bv[ext=mp4]/b[ext=mp4]",
+            "quiet": True,
+            "no_warnings": True,
+        }
+        if cookies_from_browser:
+            ydl_opts["cookiesfrombrowser"] = ("chrome",)
+
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            ydl.download(url)
+
+        for idx, (s, e) in enumerate(clips_timestamps):
+            s_time = _get_seconds(s)
+            e_time = _get_seconds(e)
+            duration = e_time - s_time
+            _ = ffmpeg.input(video_path, ss=s_time, t=duration).output(str(output_paths[idx])).run()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prefix", type=str, required=True, help="Prefix of the dataset to be downloaded")
+    parser.add_argument("--output_base", type=str, required=True, help="Base directory to save the dataset")
+    parser.add_argument("--rgb", action="store_true", help="Download RGB components of the videos")
+    parser.add_argument("--depth", action="store_true", help="Download depth components of the videos")
+
+    args = parser.parse_args()
+    output_base = Path(args.output_base)
+    output_base.mkdir(parents=True, exist_ok=True)
+
+    attributes_to_download = ["intrinsics", "pose"]
+    if args.rgb:
+        attributes_to_download.append("rgb")
+    if args.depth:
+        attributes_to_download.append("depth")
+
+    if args.prefix.startswith("dpsp"):
+        repo_id = "nvidia/vipe-dynpose-100kpp"
+    elif args.prefix.startswith("wsdg"):
+        repo_id = "nvidia/vipe-wild-sdg-1m"
+    elif args.prefix.startswith("w360"):
+        repo_id = "nvidia/vipe-web360"
+    else:
+        raise ValueError(f"Invalid prefix: {args.prefix}")
+
+    # Grab the metadata of the dataset to download
+    api = HfApi(token=os.getenv("HF_TOKEN"))
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        meta_file = api.hf_hub_download(
+            repo_id=repo_id, repo_type="dataset", filename="meta.parquet", local_dir=tmp_dir
+        )
+        df = pd.read_parquet(meta_file)
+
+    # Select rows where tar_name starts with the given prefix
+    related_videos = df.loc[df["tar_name"].str.startswith(args.prefix)]
+    related_tar_names = list(set(related_videos["tar_name"].tolist()))
+
+    print(f"Found {len(related_videos)} videos to download within {len(related_tar_names)} tar files")
+
+    for attribute in attributes_to_download:
+        for tar_name in related_tar_names:
+            remote_path = f"payload/{tar_name}/{attribute}.tar"
+            (output_base / attribute).mkdir(parents=True, exist_ok=True)
+
+            if api.file_exists(repo_id=repo_id, repo_type="dataset", filename=remote_path):
+                with tempfile.TemporaryDirectory() as tmp_dir:
+                    tar_file = api.hf_hub_download(
+                        repo_id=repo_id,
+                        repo_type="dataset",
+                        filename=remote_path,
+                        local_dir=tmp_dir,
+                    )
+                    # Extract the tar file
+                    with tarfile.open(tar_file, "r") as tar:
+                        tar.extractall(path=output_base / attribute)
+
+            else:
+                if attribute == "rgb" and args.prefix.startswith("dpsp"):
+                    # Merge videos with the same youtube_link
+                    merged_videos = {}
+                    for video_info in related_videos.iterrows():
+                        link = video_info[1]["youtube_link"]
+                        if link not in merged_videos:
+                            merged_videos[link] = []
+                        merged_videos[link].append(video_info)
+
+                    for link, merged_video in merged_videos.items():
+                        time_slices = [t[1]["youtube_timestamp"].split("-") for t in merged_video]
+                        output_links = [t[1]["sequence"] for t in merged_video]
+                        download_links = [output_base / attribute / f"{ft}.mp4" for ft in output_links]
+                        download_clips_from_url(
+                            link,
+                            time_slices,
+                            download_links,
+                            cookies_from_browser=True,
+                        )
+
+                elif attribute == "rgb" and args.prefix.startswith("w360"):
+                    # Download the depth frames from Official website
+                    gdown.download(
+                        "https://drive.google.com/file/d/1W1eLmaP16GZOeisAR1q-y9JYP9gT1CRs/view",
+                        output=str(output_base / "web360_raw.zip"),
+                        fuzzy=True,
+                        use_cookies=False,
+                    )
+                    with zipfile.ZipFile(output_base / "web360_raw.zip", "r") as zip_ref:
+                        with tempfile.TemporaryDirectory() as tmp_dir:
+                            zip_ref.extractall(tmp_dir)
+                            for video_file in Path(tmp_dir).glob("**/*.mp4"):
+                                shutil.copy(video_file, output_base / attribute / video_file.name)
+
+                else:
+                    raise ValueError(f"Attribute {attribute} is not supported for {args.prefix}")
+
+
+if __name__ == "__main__":
+    main()

third_party/vipe/scripts/vipe_to_colmap.py

@@ -0,0 +1,313 @@
+#!/usr/bin/env python3
+
+import argparse
+import logging
+
+from pathlib import Path
+from typing import Tuple
+
+import cv2
+import imageio
+import numpy as np
+import torch
+
+from scipy.spatial.transform import Rotation
+
+from vipe.slam.interface import SLAMMap
+from vipe.utils.cameras import CameraType
+from vipe.utils.depth import reliable_depth_mask_range
+from vipe.utils.io import (
+    ArtifactPath,
+    read_depth_artifacts,
+    read_intrinsics_artifacts,
+    read_pose_artifacts,
+    read_rgb_artifacts,
+)
+
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def quaternion_from_matrix(matrix: np.ndarray) -> np.ndarray:
+    """Convert rotation matrix to quaternion (w, x, y, z)."""
+    rotation = Rotation.from_matrix(matrix[:3, :3])
+    quat_xyzw = rotation.as_quat()  # Returns [x, y, z, w]
+    return np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]])  # Convert to [w, x, y, z]
+
+
+def matrix_to_colmap_pose(c2w_matrix: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Convert camera-to-world matrix to COLMAP format.
+    COLMAP uses world-to-camera transformation.
+    """
+    w2c = np.linalg.inv(c2w_matrix)
+    quaternion = quaternion_from_matrix(w2c)
+    translation = w2c[:3, 3]
+    return quaternion, translation
+
+
+def write_cameras_txt(output_dir: Path, artifact: ArtifactPath, frame_width: int, frame_height: int):
+    """Write COLMAP cameras.txt file."""
+    cameras_file = output_dir / "cameras.txt"
+
+    _, intrinsics, camera_types = read_intrinsics_artifacts(artifact.intrinsics_path)
+
+    # Use first frame's intrinsics (assuming constant intrinsics)
+    assert camera_types[0] == CameraType.PINHOLE, "Only PINHOLE camera type is supported"
+    fx, fy, cx, cy = intrinsics[0].cpu().numpy()
+
+    with open(cameras_file, "w") as f:
+        f.write("# Camera list with one line of data per camera:\n")
+        f.write("#   CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n")
+        f.write("# Number of cameras: 1\n")
+
+        fx, fy, cx, cy = intrinsics[0]
+
+        # COLMAP camera format: CAMERA_ID MODEL WIDTH HEIGHT fx fy cx cy
+        f.write(f"1 PINHOLE {frame_width} {frame_height} {fx:.6f} {fy:.6f} {cx:.6f} {cy:.6f}\n")
+
+    logger.info(f"Written cameras.txt with intrinsics: fx={fx:.2f}, fy={fy:.2f}, cx={cx:.2f}, cy={cy:.2f}")
+
+
+def write_images_txt(output_dir: Path, artifact: ArtifactPath):
+    """Write COLMAP images.txt file."""
+    images_file = output_dir / "images.txt"
+
+    # Load pose data
+    pose_data = np.load(artifact.pose_path)
+    poses = pose_data["data"]  # Shape: (N, 4, 4)
+    indices = pose_data["inds"]  # Frame indices
+
+    with open(images_file, "w") as f:
+        f.write("# Image list with two lines of data per image:\n")
+        f.write("#   IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n")
+        f.write("#   POINTS2D[] as (X, Y, POINT3D_ID)\n")
+        f.write(f"# Number of images: {len(poses)}\n")
+
+        for i, (pose_matrix, frame_idx) in enumerate(zip(poses, indices)):
+            # Convert pose to COLMAP format
+            quaternion, translation = matrix_to_colmap_pose(pose_matrix)
+            qw, qx, qy, qz = quaternion
+            tx, ty, tz = translation
+
+            # Image filename
+            image_name = f"images/frame_{frame_idx:06d}.jpg"
+
+            # Write image line
+            f.write(f"{i + 1} {qw:.9f} {qx:.9f} {qy:.9f} {qz:.9f} {tx:.9f} {ty:.9f} {tz:.9f} 1 {image_name}\n")
+            # Empty points2D line (no 2D-3D correspondences)
+            f.write("\n")
+
+    logger.info(f"Written images.txt with {len(poses)} images")
+
+
+def write_points3d_txt_from_slam_map(output_dir: Path, artifact: ArtifactPath):
+    """Write points3D.txt from SLAM map (placeholder implementation)."""
+    assert artifact.slam_map_path.exists(), "SLAM map not found, please refer to README.md for more details."
+
+    slam_map = SLAMMap.load(artifact.slam_map_path, device=torch.device("cpu"))
+
+    points3d_file = output_dir / "points3D.txt"
+    with open(points3d_file, "w") as f:
+        f.write("# 3D point list with one line of data per point:\n")
+        f.write("#   POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n")
+        f.write(f"# Number of points: {slam_map.dense_disp_xyz.shape[0]}\n")
+
+        point_id = 1
+        for keyframe_idx, frame_idx in enumerate(slam_map.dense_disp_frame_inds):
+            xyz, rgb = slam_map.get_dense_disp_pcd(keyframe_idx)
+            xyz = xyz.cpu().numpy()
+            rgb = rgb.cpu().numpy()
+
+            for xyz, rgb in zip(xyz, rgb):
+                x, y, z = xyz
+                r, g, b = (rgb * 255).astype(np.uint8)
+                f.write(f"{point_id} {x:.6f} {y:.6f} {z:.6f} {r} {g} {b} 0.0 {frame_idx} {point_id} 0 0 0 0\n")
+                point_id += 1
+
+
+def write_points3d_txt_from_depth(
+    output_dir: Path, artifact: ArtifactPath, depth_step: int, spatial_subsample: int = 4
+):
+    """Write empty COLMAP points3D.txt file."""
+    _, pose_data = read_pose_artifacts(artifact.pose_path)
+    _, intrinsics, camera_types = read_intrinsics_artifacts(artifact.intrinsics_path)
+    camera_type = camera_types[0]
+    points3d_file = output_dir / "points3D.txt"
+
+    image_dir = output_dir / "images"
+    images = sorted(list(image_dir.glob("*.jpg")))
+    # Collect all 3D points first
+    all_points = []
+    point_id = 1
+
+    rays: np.ndarray | None = None
+
+    for idx, (_, depth) in enumerate(read_depth_artifacts(artifact.depth_path)):
+        if idx % 30 == 0:
+            logger.info(f"Processed {idx} depth maps")
+
+        if idx % depth_step != 0:
+            continue
+
+        rgb = cv2.cvtColor(cv2.imread(str(images[idx]), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
+        frame_height, frame_width = rgb.shape[:2]
+        rgb = rgb[::spatial_subsample, ::spatial_subsample]
+
+        if rays is None:
+            camera_model = camera_type.build_camera_model(intrinsics[idx])
+            disp_v, disp_u = torch.meshgrid(
+                torch.arange(frame_height).float()[::spatial_subsample],
+                torch.arange(frame_width).float()[::spatial_subsample],
+                indexing="ij",
+            )
+            if camera_type == CameraType.PANORAMA:
+                disp_v = disp_v / (frame_height - 1)
+                disp_u = disp_u / (frame_width - 1)
+            disp = torch.ones_like(disp_v)
+            pts, _, _ = camera_model.iproj_disp(disp, disp_u, disp_v)
+            rays = pts[..., :3].numpy()
+            if camera_type != CameraType.PANORAMA:
+                rays /= rays[..., 2:3]
+
+        if depth is not None:
+            pcd = rays * depth.numpy()[::spatial_subsample, ::spatial_subsample, None]
+            depth_mask = reliable_depth_mask_range(depth)[::spatial_subsample, ::spatial_subsample].numpy()
+            rgb, pcd = rgb[depth_mask], pcd[depth_mask]
+            c2w_matrix = pose_data[idx].matrix().numpy()
+            pcd = pcd @ c2w_matrix[:3, :3].T + c2w_matrix[:3, 3][None]
+
+            for pts_rgb, pts_xyz in zip(rgb, pcd):
+                all_points.append(
+                    (
+                        point_id,
+                        pts_xyz[0],
+                        pts_xyz[1],
+                        pts_xyz[2],
+                        int(pts_rgb[0]),
+                        int(pts_rgb[1]),
+                        int(pts_rgb[2]),
+                        0.0,
+                        idx + 1,
+                    )
+                )
+                point_id += 1
+
+    # Write points to file
+    with open(points3d_file, "w") as f:
+        f.write("# 3D point list with one line of data per point:\n")
+        f.write("#   POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n")
+        f.write(f"# Number of points: {len(all_points)}\n")
+
+        for point_id, point_data in enumerate(all_points):
+            point_id, x, y, z, r, g, b, error, image_id = point_data
+            # The last 4 values are for visualization purposes.
+            f.write(f"{point_id} {x:.6f} {y:.6f} {z:.6f} {r} {g} {b} {error:.6f} {image_id} {point_id} 0 0 0 0\n")
+
+    logger.info(f"Written points3D.txt with {len(all_points)} points")
+
+
+def extract_frames(artifact: ArtifactPath, output_dir: Path) -> Tuple[int, int]:
+    """Extract frames from video to individual image files."""
+    video_path = artifact.rgb_path
+    images_dir = output_dir / "images"
+    images_dir.mkdir(exist_ok=True)
+
+    logger.info(f"Extracting frames from {video_path}")
+
+    for frame_idx, rgb in read_rgb_artifacts(video_path):
+        frame_path = images_dir / f"frame_{frame_idx:06d}.jpg"
+        frame_height, frame_width = rgb.shape[:2]
+        imageio.imwrite(str(frame_path), (rgb.cpu().numpy() * 255).astype(np.uint8))
+        if frame_idx % 30 == 0:
+            logger.info(f"Extracted {frame_idx} frames")
+
+    logger.info(f"Extracted {frame_idx} frames to {images_dir}")
+
+    return frame_width, frame_height
+
+
+def convert_vipe_to_colmap(artifact: ArtifactPath, output_path: Path, depth_step: int, use_slam_map: bool):
+    """Convert ViPE reconstruction results to COLMAP format."""
+
+    logger.info(
+        f"Converting ViPE results from {artifact.base_path} ({artifact.artifact_name}) to COLMAP format at {output_path}"
+    )
+
+    # Verify required files exist
+    required_files = [artifact.rgb_path, artifact.pose_path, artifact.intrinsics_path, artifact.depth_path]
+    for file_path in required_files:
+        if not file_path.exists():
+            raise FileNotFoundError(f"Required file not found: {file_path}")
+
+    # Create output directory
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    # Extract frames and get video dimensions
+    frame_width, frame_height = extract_frames(artifact, output_path)
+
+    # Write COLMAP files
+    write_cameras_txt(output_path, artifact, frame_width, frame_height)
+    write_images_txt(output_path, artifact)
+    if use_slam_map:
+        write_points3d_txt_from_slam_map(output_path, artifact)
+    else:
+        write_points3d_txt_from_depth(output_path, artifact, depth_step)
+
+    logger.info("COLMAP conversion completed successfully!")
+    logger.info(f"Output directory: {output_path}")
+    logger.info("Files created:")
+    logger.info("  - cameras.txt: Camera intrinsics")
+    logger.info("  - images.txt: Camera poses")
+    logger.info("  - points3D.txt: 3D points")
+    logger.info("  - images/: Individual frame images")
+
+
+def main():
+    """Main function for ViPE to COLMAP conversion script."""
+    parser = argparse.ArgumentParser(
+        description="Convert ViPE reconstruction results to COLMAP format",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("vipe_path", type=Path, help="Path to ViPE results directory")
+    parser.add_argument(
+        "--sequence",
+        "-s",
+        type=str,
+        help="Sequence name (if not provided will convert all sequences in the directory)",
+        default=None,
+    )
+    parser.add_argument("--use_slam_map", action="store_true", help="Use SLAM map to unproject depth maps")
+    parser.add_argument("--depth_step", type=int, default=16, help="Step size for depth extraction (default: 16)")
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=Path,
+        default=None,
+        help="Output directory for COLMAP format (default: <vipe_path>_colmap)",
+    )
+
+    args = parser.parse_args()
+
+    if not args.vipe_path.exists():
+        print(f"Error: ViPE path '{args.vipe_path}' does not exist.")
+        return 1
+
+    # Find artifacts
+    artifacts = list(ArtifactPath.glob_artifacts(args.vipe_path, use_video=True))
+    if args.sequence is not None:
+        artifacts = [artifact for artifact in artifacts if artifact.artifact_name == args.sequence]
+
+    # Set default output path
+    if args.output is None:
+        args.output = args.vipe_path.parent / f"{args.vipe_path.name}_colmap"
+
+    for artifact in artifacts:
+        convert_vipe_to_colmap(artifact, args.output / artifact.artifact_name, args.depth_step, args.use_slam_map)
+    return 0
+
+
+if __name__ == "__main__":
+    exit(main())

third_party/vipe/setup.py

@@ -0,0 +1,83 @@
+import os
+import re
+import shutil
+import tarfile
+import tempfile
+
+from setuptools import find_packages, setup
+from urllib.request import urlretrieve
+
+try:
+    import torch
+    import torch.version
+
+    from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+    torch_version = torch.version.__version__.split(".")[:2]
+    cuda_version = torch.version.cuda
+
+    # This will be e.g. "+pt23cu121"
+    assert cuda_version is not None, "Pytorch CUDA is required for this installation."
+    version_suffix = f"+pt{torch_version[0]}{torch_version[1]}cu{cuda_version.replace('.', '')}"
+
+except ImportError:
+    raise ValueError("Pytorch not found, please install it first.")
+
+PACKAGE_NAME = "vipe"
+
+# Avoid directly importing the package
+with open(f"{PACKAGE_NAME}/__init__.py", "r") as fh:
+    __version__ = re.findall(r"__version__ = \"(.*?)\"", fh.read())[0]
+__version__ += version_suffix
+
+coder_finder_path = f"{PACKAGE_NAME}/ext/specs.py"
+code_finder_namespace = {"__file__": coder_finder_path}
+with open(coder_finder_path, "r") as fh:
+    exec(fh.read(), code_finder_namespace)
+get_sources = code_finder_namespace["get_sources"]
+get_cpp_flags = code_finder_namespace["get_cpp_flags"]
+get_cuda_flags = code_finder_namespace["get_cuda_flags"]
+
+# Setup CUDA_HOME for conda environment for consistency
+if "CONDA_PREFIX" in os.environ:
+    conda_nvcc_path = os.path.join(os.environ["CONDA_PREFIX"], "bin", "nvcc")
+    if os.path.exists(conda_nvcc_path):
+        os.environ["PYTORCH_NVCC"] = conda_nvcc_path
+
+# Download the put Eigen 3.4 in a correct place
+cpp_flags = get_cpp_flags()
+cuda_flags = get_cuda_flags()
+if os.environ.get("USE_SYSTEM_EIGEN", "0") == "0":
+    eigen_include_dir = "csrc/include/eigen3"
+    eigen_url = "https://gitlab.com/libeigen/eigen/-/archive/3.4.0/eigen-3.4.0.tar.gz"
+
+    if not os.path.exists(eigen_include_dir):
+        os.makedirs(eigen_include_dir, exist_ok=True)
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            tmp_tar_path = os.path.join(temp_dir, "eigen.gz")
+            extracted_dir = os.path.join(temp_dir, "eigen-extracted")
+            urlretrieve(eigen_url, tmp_tar_path)
+            with tarfile.open(tmp_tar_path, "r:gz") as tar:
+                tar.extractall(path=extracted_dir)
+
+            shutil.move(os.path.join(extracted_dir, "eigen-3.4.0", "Eigen"), eigen_include_dir)
+
+    # Use full path
+    additional_include_path = os.path.join(os.path.dirname(__file__), "csrc/include")
+    cpp_flags += ["-isystem", additional_include_path]
+    cuda_flags += ["-isystem", additional_include_path]
+
+packages = find_packages()
+setup(
+    packages=packages,
+    version=__version__,
+    ext_modules=[
+        CUDAExtension(
+            f"{PACKAGE_NAME}_ext",
+            sources=get_sources(),  # type: ignore
+            extra_compile_args={"cxx": cpp_flags, "nvcc": cuda_flags},  # type: ignore
+        )
+    ],
+    cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},
+)

third_party/vipe/vipe/__init__.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+from omegaconf import OmegaConf
+
+from vipe.pipeline import make_pipeline
+
+
+__version__ = "0.1.1"
+__version_info__ = (0, 1, 1)
+
+if not OmegaConf.has_resolver("eq"):
+    OmegaConf.register_new_resolver("eq", lambda a, b: a == b)
+if not OmegaConf.has_resolver("neq"):
+    OmegaConf.register_new_resolver("neq", lambda a, b: a != b)
+
+
+def get_config_path() -> Path:
+    return Path(__file__).parent.parent / "configs"

third_party/vipe/vipe/cli/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

third_party/vipe/vipe/cli/main.py

@@ -0,0 +1,115 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+import click
+import hydra
+
+from vipe import get_config_path, make_pipeline
+from vipe.streams.base import ProcessedVideoStream
+from vipe.streams.raw_mp4_stream import RawMp4Stream
+from vipe.streams.frame_dir_stream import FrameDirStream
+from vipe.utils.logging import configure_logging
+from vipe.utils.viser import run_viser
+
+
+@click.command()
+@click.argument("video", type=click.Path(exists=True, path_type=Path), required=False)
+@click.option(
+    "--image-dir",
+    type=click.Path(exists=True, path_type=Path),
+    help="Directory containing image frames",
+)
+@click.option(
+    "--output",
+    "-o",
+    type=click.Path(path_type=Path),
+    help="Output directory (default: current directory)",
+    default=Path.cwd() / "vipe_results",
+)
+@click.option("--pipeline", "-p", default="default", help="Pipeline configuration to use (default: 'default')")
+@click.option("--visualize", "-v", is_flag=True, help="Enable visualization of intermediate results")
+def infer(video: Path, image_dir: Path, output: Path, pipeline: str, visualize: bool):
+    """Run inference on a video file or directory of images."""
+
+    logger = configure_logging()
+
+    # Validate that exactly one input source is provided
+    if not video and not image_dir:
+        click.echo("Error: Must provide either a video file or --image-dir", err=True)
+        raise click.Abort()
+    
+    if video and image_dir:
+        click.echo("Error: Cannot provide both video file and --image-dir", err=True)
+        raise click.Abort()
+
+    overrides = [f"pipeline={pipeline}", f"pipeline.output.path={output}", "pipeline.output.save_artifacts=true"]
+    if visualize:
+        overrides.append("pipeline.output.save_viz=true")
+        overrides.append("pipeline.slam.visualize=true")
+    else:
+        overrides.append("pipeline.output.save_viz=false")
+
+    # Set up stream configuration based on input type
+    if image_dir:
+        overrides.extend([
+            "streams=frame_dir_stream",
+            f"streams.base_path={image_dir}"
+        ])
+        input_path = image_dir
+        input_desc = f"image directory {image_dir}"
+    else:
+        input_path = video
+        input_desc = f"video {video}"
+
+    with hydra.initialize_config_dir(config_dir=str(get_config_path()), version_base=None):
+        args = hydra.compose("default", overrides=overrides)
+
+    logger.info(f"Processing {input_desc}...")
+    vipe_pipeline = make_pipeline(args.pipeline)
+
+    if image_dir:
+        # Use frame directory stream
+        video_stream = ProcessedVideoStream(FrameDirStream(image_dir), []).cache(desc="Reading image frames")
+    else:
+        # Some input videos can be malformed, so we need to cache the videos to obtain correct number of frames.
+        video_stream = ProcessedVideoStream(RawMp4Stream(video), []).cache(desc="Reading video stream")
+
+    vipe_pipeline.run(video_stream)
+    logger.info("Finished")
+
+
+@click.command()
+@click.argument("data_path", type=click.Path(exists=True, path_type=Path), default=Path.cwd() / "vipe_results")
+@click.option("--port", "-p", default=20540, type=int, help="Port for the visualization server (default: 20540)")
+def visualize(data_path: Path, port: int):
+    run_viser(data_path, port)
+
+
+@click.group()
+@click.version_option()
+def main():
+    """NVIDIA Video Pose Engine (ViPE) CLI"""
+    pass
+
+
+# Add subcommands
+main.add_command(infer)
+main.add_command(visualize)
+
+
+if __name__ == "__main__":
+    main()

third_party/vipe/vipe/ext/__init__.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from torch.utils.cpp_extension import load
+
+from vipe.ext.specs import get_cpp_flags, get_cuda_flags, get_sources
+
+
+try:
+    import vipe_ext as _C
+
+    vipe_ext_not_found = False
+except ImportError:
+    vipe_ext_not_found = True
+
+if vipe_ext_not_found or os.environ.get("VIPE_EXT_JIT", "0") == "1":
+    _C = load(
+        name="vipe_ext_jit",
+        sources=get_sources(),
+        extra_cflags=get_cpp_flags(),
+        extra_cuda_cflags=get_cuda_flags(),
+        verbose=True,
+    )
+
+# Reference to submodules
+droid_net_ext = _C.droid_net_ext
+grounding_dino_ext = _C.grounding_dino_ext
+utils_ext = _C.utils_ext
+slam_ext = _C.slam_ext
+scatter_ext = _C.scatter_ext
+lietorch_ext = _C.lietorch_ext
+corr_ext = _C.corr_ext

third_party/vipe/vipe/ext/corr/__init__.py

@@ -0,0 +1,5 @@
+# This file includes code originally from the Pytorch Correlation repository:
+# https://github.com/ClementPinard/Pytorch-Correlation-extension
+# Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+
+from .spatial_correlation_sampler import SpatialCorrelationSampler, spatial_correlation_sample

third_party/vipe/vipe/ext/corr/spatial_correlation_sampler.py

@@ -0,0 +1,126 @@
+# This file includes code originally from the Pytorch Correlation repository:
+# https://github.com/ClementPinard/Pytorch-Correlation-extension
+# Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+
+from torch import nn
+from torch.autograd import Function
+from torch.autograd.function import once_differentiable
+from torch.nn.modules.utils import _pair
+
+from vipe.ext import corr_ext as correlation
+
+
+def spatial_correlation_sample(
+    input1, input2, kernel_size=1, patch_size=1, stride=1, padding=0, dilation=1, dilation_patch=1
+):
+    """Apply spatial correlation sampling on from input1 to input2,
+
+    Every parameter except input1 and input2 can be either single int
+    or a pair of int. For more information about Spatial Correlation
+    Sampling, see this page.
+    https://lmb.informatik.uni-freiburg.de/Publications/2015/DFIB15/
+
+    Args:
+        input1 : The first parameter.
+        input2 : The second parameter.
+        kernel_size : total size of your correlation kernel, in pixels
+        patch_size : total size of your patch, determining how many
+            different shifts will be applied
+        stride : stride of the spatial sampler, will modify output
+            height and width
+        padding : padding applied to input1 and input2 before applying
+            the correlation sampling, will modify output height and width
+        dilation_patch : step for every shift in patch
+
+    Returns:
+        Tensor: Result of correlation sampling
+
+    """
+    return SpatialCorrelationSamplerFunction.apply(
+        input1, input2, kernel_size, patch_size, stride, padding, dilation, dilation_patch
+    )
+
+
+class SpatialCorrelationSamplerFunction(Function):
+    @staticmethod
+    def forward(ctx, input1, input2, kernel_size=1, patch_size=1, stride=1, padding=0, dilation=1, dilation_patch=1):
+        ctx.save_for_backward(input1, input2)
+        kH, kW = ctx.kernel_size = _pair(kernel_size)
+        patchH, patchW = ctx.patch_size = _pair(patch_size)
+        padH, padW = ctx.padding = _pair(padding)
+        dilationH, dilationW = ctx.dilation = _pair(dilation)
+        dilation_patchH, dilation_patchW = ctx.dilation_patch = _pair(dilation_patch)
+        dH, dW = ctx.stride = _pair(stride)
+
+        output = correlation.forward(
+            input1,
+            input2,
+            kH,
+            kW,
+            patchH,
+            patchW,
+            padH,
+            padW,
+            dilationH,
+            dilationW,
+            dilation_patchH,
+            dilation_patchW,
+            dH,
+            dW,
+        )
+
+        return output
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        input1, input2 = ctx.saved_tensors
+
+        kH, kW = ctx.kernel_size
+        patchH, patchW = ctx.patch_size
+        padH, padW = ctx.padding
+        dilationH, dilationW = ctx.dilation
+        dilation_patchH, dilation_patchW = ctx.dilation_patch
+        dH, dW = ctx.stride
+
+        grad_input1, grad_input2 = correlation.backward(
+            input1,
+            input2,
+            grad_output,
+            kH,
+            kW,
+            patchH,
+            patchW,
+            padH,
+            padW,
+            dilationH,
+            dilationW,
+            dilation_patchH,
+            dilation_patchW,
+            dH,
+            dW,
+        )
+        return grad_input1, grad_input2, None, None, None, None, None, None
+
+
+class SpatialCorrelationSampler(nn.Module):
+    def __init__(self, kernel_size=1, patch_size=1, stride=1, padding=0, dilation=1, dilation_patch=1):
+        super(SpatialCorrelationSampler, self).__init__()
+        self.kernel_size = kernel_size
+        self.patch_size = patch_size
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.dilation_patch = dilation_patch
+
+    def forward(self, input1, input2):
+        return SpatialCorrelationSamplerFunction.apply(
+            input1,
+            input2,
+            self.kernel_size,
+            self.patch_size,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.dilation_patch,
+        )

third_party/vipe/vipe/ext/lietorch/__init__.py

@@ -0,0 +1,6 @@
+# This file includes code originally from the lietorch repository:
+# https://github.com/princeton-vl/lietorch
+# Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+
+__all__ = ["groups"]
+from .groups import SE3, SO3, LieGroupParameter, RxSO3, Sim3, cat, stack

third_party/vipe/vipe/ext/lietorch/broadcasting.py

@@ -0,0 +1,37 @@
+# This file includes code originally from the lietorch repository:
+# https://github.com/princeton-vl/lietorch
+# Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+
+import numpy as np
+import torch
+
+
+def check_broadcastable(x, y):
+    assert len(x.shape) == len(y.shape)
+    for n, m in zip(x.shape[:-1], y.shape[:-1]):
+        assert n == m or n == 1 or m == 1
+
+
+def broadcast_inputs(x, y):
+    """Automatic broadcasting of missing dimensions"""
+    if y is None:
+        xs, xd = x.shape[:-1], x.shape[-1]
+        return (x.view(-1, xd).contiguous(),), x.shape[:-1]
+
+    check_broadcastable(x, y)
+
+    xs, xd = x.shape[:-1], x.shape[-1]
+    ys, yd = y.shape[:-1], y.shape[-1]
+    out_shape = [max(n, m) for (n, m) in zip(xs, ys)]
+
+    if x.shape[:-1] == y.shape[-1]:
+        x1 = x.view(-1, xd)
+        y1 = y.view(-1, yd)
+
+    else:
+        x_expand = [m if n == 1 else 1 for (n, m) in zip(xs, ys)]
+        y_expand = [n if m == 1 else 1 for (n, m) in zip(xs, ys)]
+        x1 = x.repeat(x_expand + [1]).reshape(-1, xd).contiguous()
+        y1 = y.repeat(y_expand + [1]).reshape(-1, yd).contiguous()
+
+    return (x1, y1), tuple(out_shape)

third_party/vipe/vipe/ext/lietorch/group_ops.py

@@ -0,0 +1,124 @@
+# This file includes code originally from the lietorch repository:
+# https://github.com/princeton-vl/lietorch
+# Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+
+import torch
+import torch.nn.functional as F
+
+from vipe.ext import lietorch_ext as lietorch_backends
+
+
+class GroupOp(torch.autograd.Function):
+    """group operation base class"""
+
+    @classmethod
+    def forward(cls, ctx, group_id, *inputs):
+        ctx.group_id = group_id
+        ctx.save_for_backward(*inputs)
+        out = cls.forward_op(ctx.group_id, *inputs)
+        return out
+
+    @classmethod
+    def backward(cls, ctx, grad):
+        error_str = "Backward operation not implemented for {}".format(cls)
+        assert cls.backward_op is not None, error_str
+
+        inputs = ctx.saved_tensors
+        grad = grad.contiguous()
+        grad_inputs = cls.backward_op(ctx.group_id, grad, *inputs)
+        return (None,) + tuple(grad_inputs)
+
+
+class Exp(GroupOp):
+    """exponential map"""
+
+    forward_op, backward_op = lietorch_backends.expm, lietorch_backends.expm_backward
+
+
+class Log(GroupOp):
+    """logarithm map"""
+
+    forward_op, backward_op = lietorch_backends.logm, lietorch_backends.logm_backward
+
+
+class Inv(GroupOp):
+    """group inverse"""
+
+    forward_op, backward_op = lietorch_backends.inv, lietorch_backends.inv_backward
+
+
+class Mul(GroupOp):
+    """group multiplication"""
+
+    forward_op, backward_op = lietorch_backends.mul, lietorch_backends.mul_backward
+
+
+class Adj(GroupOp):
+    """adjoint operator"""
+
+    forward_op, backward_op = lietorch_backends.adj, lietorch_backends.adj_backward
+
+
+class AdjT(GroupOp):
+    """adjoint operator"""
+
+    forward_op, backward_op = lietorch_backends.adjT, lietorch_backends.adjT_backward
+
+
+class Act3(GroupOp):
+    """action on point"""
+
+    forward_op, backward_op = lietorch_backends.act, lietorch_backends.act_backward
+
+
+class Act4(GroupOp):
+    """action on point"""
+
+    forward_op, backward_op = lietorch_backends.act4, lietorch_backends.act4_backward
+
+
+class Jinv(GroupOp):
+    """adjoint operator"""
+
+    forward_op, backward_op = lietorch_backends.Jinv, None
+
+
+class ToMatrix(GroupOp):
+    """convert to matrix representation"""
+
+    forward_op, backward_op = lietorch_backends.as_matrix, None
+
+
+### conversion operations to/from Euclidean embeddings ###
+
+
+class FromVec(torch.autograd.Function):
+    """convert vector into group object"""
+
+    @classmethod
+    def forward(cls, ctx, group_id, *inputs):
+        ctx.group_id = group_id
+        ctx.save_for_backward(*inputs)
+        return inputs[0]
+
+    @classmethod
+    def backward(cls, ctx, grad):
+        inputs = ctx.saved_tensors
+        J = lietorch_backends.projector(ctx.group_id, *inputs)
+        return None, torch.matmul(grad.unsqueeze(-2), torch.linalg.pinv(J)).squeeze(-2)
+
+
+class ToVec(torch.autograd.Function):
+    """convert group object to vector"""
+
+    @classmethod
+    def forward(cls, ctx, group_id, *inputs):
+        ctx.group_id = group_id
+        ctx.save_for_backward(*inputs)
+        return inputs[0]
+
+    @classmethod
+    def backward(cls, ctx, grad):
+        inputs = ctx.saved_tensors
+        J = lietorch_backends.projector(ctx.group_id, *inputs)
+        return None, torch.matmul(grad.unsqueeze(-2), J).squeeze(-2)

third_party/vipe/vipe/ext/lietorch/groups.py

@@ -0,0 +1,328 @@
+# This file includes code originally from the lietorch repository:
+# https://github.com/princeton-vl/lietorch
+# Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+
+import numpy as np
+import torch
+
+from .broadcasting import broadcast_inputs
+
+# group operations implemented in cuda
+from .group_ops import Act3, Act4, Adj, AdjT, Exp, FromVec, Inv, Jinv, Log, Mul, ToMatrix, ToVec
+
+
+class LieGroupParameter(torch.Tensor):
+    """Wrapper class for LieGroup"""
+
+    from torch._C import _disabled_torch_function_impl
+
+    __torch_function__ = _disabled_torch_function_impl
+
+    def __new__(cls, group, requires_grad=True):
+        data = torch.zeros(group.tangent_shape, device=group.data.device, dtype=group.data.dtype, requires_grad=True)
+
+        return torch.Tensor._make_subclass(cls, data, requires_grad)
+
+    def __init__(self, group):
+        self.group = group
+
+    def retr(self):
+        return self.group.retr(self)
+
+    def log(self):
+        return self.retr().log()
+
+    def inv(self):
+        return self.retr().inv()
+
+    def adj(self, a):
+        return self.retr().adj(a)
+
+    def __mul__(self, other):
+        if isinstance(other, LieGroupParameter):
+            return self.retr() * other.retr()
+        else:
+            return self.retr() * other
+
+    def add_(self, update, alpha):
+        self.group = self.group.exp(alpha * update) * self.group
+
+    def __getitem__(self, index):
+        return self.retr().__getitem__(index)
+
+
+class LieGroup:
+    """Base class for Lie Group"""
+
+    def __init__(self, data):
+        self.data = data
+
+    def __repr__(self):
+        return "{}: size={}, device={}, dtype={}".format(self.group_name, self.shape, self.device, self.dtype)
+
+    @property
+    def shape(self):
+        return self.data.shape[:-1]
+
+    @property
+    def device(self):
+        return self.data.device
+
+    @property
+    def dtype(self):
+        return self.data.dtype
+
+    def vec(self):
+        return self.apply_op(ToVec, self.data)
+
+    @property
+    def tangent_shape(self):
+        return self.data.shape[:-1] + (self.manifold_dim,)
+
+    @classmethod
+    def Identity(cls, *batch_shape, **kwargs):
+        """Construct identity element with batch shape"""
+
+        if isinstance(batch_shape[0], tuple):
+            batch_shape = batch_shape[0]
+
+        elif isinstance(batch_shape[0], list):
+            batch_shape = tuple(batch_shape[0])
+
+        numel = np.prod(batch_shape)
+        data = cls.id_elem.reshape(1, -1)
+
+        if "device" in kwargs:
+            data = data.to(kwargs["device"])
+
+        if "dtype" in kwargs:
+            data = data.type(kwargs["dtype"])
+
+        data = data.repeat(numel, 1)
+        return cls(data).view(batch_shape)
+
+    @classmethod
+    def IdentityLike(cls, G):
+        return cls.Identity(G.shape, device=G.data.device, dtype=G.data.dtype)
+
+    @classmethod
+    def InitFromVec(cls, data):
+        return cls(cls.apply_op(FromVec, data))
+
+    @classmethod
+    def Random(cls, *batch_shape, sigma=1.0, **kwargs):
+        """Construct random element with batch_shape by random sampling in tangent space"""
+
+        if isinstance(batch_shape[0], tuple):
+            batch_shape = batch_shape[0]
+
+        elif isinstance(batch_shape[0], list):
+            batch_shape = tuple(batch_shape[0])
+
+        tangent_shape = batch_shape + (cls.manifold_dim,)
+        xi = torch.randn(tangent_shape, **kwargs)
+        return cls.exp(sigma * xi)
+
+    @classmethod
+    def apply_op(cls, op, x, y=None):
+        """Apply group operator"""
+        inputs, out_shape = broadcast_inputs(x, y)
+
+        data = op.apply(cls.group_id, *inputs)
+        return data.view(out_shape + (-1,))
+
+    @classmethod
+    def exp(cls, x):
+        """exponential map: x -> X"""
+        return cls(cls.apply_op(Exp, x))
+
+    def quaternion(self):
+        """extract quaternion"""
+        return self.apply_op(Quat, self.data)
+
+    def log(self):
+        """logarithm map"""
+        return self.apply_op(Log, self.data)
+
+    def inv(self):
+        """group inverse"""
+        return self.__class__(self.apply_op(Inv, self.data))
+
+    def mul(self, other):
+        """group multiplication"""
+        return self.__class__(self.apply_op(Mul, self.data, other.data))
+
+    def retr(self, a):
+        """retraction: Exp(a) * X"""
+        dX = self.__class__.apply_op(Exp, a)
+        return self.__class__(self.apply_op(Mul, dX, self.data))
+
+    def adj(self, a):
+        """adjoint operator: b = A(X) * a"""
+        return self.apply_op(Adj, self.data, a)
+
+    def adjT(self, a):
+        """transposed adjoint operator: b = a * A(X)"""
+        return self.apply_op(AdjT, self.data, a)
+
+    def Jinv(self, a):
+        return self.apply_op(Jinv, self.data, a)
+
+    def act(self, p):
+        """action on a point cloud"""
+
+        # action on point
+        if p.shape[-1] == 3:
+            return self.apply_op(Act3, self.data, p)
+
+        # action on homogeneous point
+        elif p.shape[-1] == 4:
+            return self.apply_op(Act4, self.data, p)
+
+    def matrix(self):
+        """convert element to 4x4 matrix"""
+        I = torch.eye(4, dtype=self.dtype, device=self.device)
+        I = I.view([1] * (len(self.data.shape) - 1) + [4, 4])
+        return self.__class__(self.data[..., None, :]).act(I).transpose(-1, -2)
+
+    def translation(self):
+        """extract translation component"""
+        p = torch.as_tensor([0.0, 0.0, 0.0, 1.0], dtype=self.dtype, device=self.device)
+        p = p.view(
+            [1] * (len(self.data.shape) - 1)
+            + [
+                4,
+            ]
+        )
+        return self.apply_op(Act4, self.data, p)
+
+    def detach(self):
+        return self.__class__(self.data.detach())
+
+    def view(self, dims):
+        data_reshaped = self.data.view(dims + (self.embedded_dim,))
+        return self.__class__(data_reshaped)
+
+    def __mul__(self, other):
+        # group multiplication
+        if isinstance(other, LieGroup):
+            return self.mul(other)
+
+        # action on point
+        elif isinstance(other, torch.Tensor):
+            return self.act(other)
+
+    def __getitem__(self, index):
+        return self.__class__(self.data[index])
+
+    def __setitem__(self, index, item):
+        self.data[index] = item.data
+
+    def to(self, *args, **kwargs):
+        return self.__class__(self.data.to(*args, **kwargs))
+
+    def cpu(self):
+        return self.__class__(self.data.cpu())
+
+    def cuda(self):
+        return self.__class__(self.data.cuda())
+
+    def float(self, device):
+        return self.__class__(self.data.float())
+
+    def double(self, device):
+        return self.__class__(self.data.double())
+
+    def unbind(self, dim=0):
+        return [self.__class__(x) for x in self.data.unbind(dim=dim)]
+
+
+class SO3(LieGroup):
+    group_name = "SO3"
+    group_id = 1
+    manifold_dim = 3
+    embedded_dim = 4
+
+    # unit quaternion
+    id_elem = torch.as_tensor([0.0, 0.0, 0.0, 1.0])
+
+    def __init__(self, data):
+        if isinstance(data, SE3):
+            data = data.data[..., 3:7]
+
+        super(SO3, self).__init__(data)
+
+
+class RxSO3(LieGroup):
+    group_name = "RxSO3"
+    group_id = 2
+    manifold_dim = 4
+    embedded_dim = 5
+
+    # unit quaternion
+    id_elem = torch.as_tensor([0.0, 0.0, 0.0, 1.0, 1.0])
+
+    def __init__(self, data):
+        if isinstance(data, Sim3):
+            data = data.data[..., 3:8]
+
+        super(RxSO3, self).__init__(data)
+
+
+class SE3(LieGroup):
+    group_name = "SE3"
+    group_id = 3
+    manifold_dim = 6
+    embedded_dim = 7
+
+    # translation, unit quaternion
+    id_elem = torch.as_tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0])
+
+    def __init__(self, data):
+        if isinstance(data, SO3):
+            translation = torch.zeros_like(data.data[..., :3])
+            data = torch.cat([translation, data.data], -1)
+
+        super(SE3, self).__init__(data)
+
+    def scale(self, s):
+        t, q = self.data.split([3, 4], -1)
+        t = t * s.unsqueeze(-1)
+        return SE3(torch.cat([t, q], dim=-1))
+
+
+class Sim3(LieGroup):
+    group_name = "Sim3"
+    group_id = 4
+    manifold_dim = 7
+    embedded_dim = 8
+
+    # translation, unit quaternion, scale
+    id_elem = torch.as_tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0])
+
+    def __init__(self, data):
+        if isinstance(data, SO3):
+            scale = torch.ones_like(SO3.data[..., :1])
+            translation = torch.zeros_like(SO3.data[..., :3])
+            data = torch.cat([translation, SO3.data, scale], -1)
+
+        elif isinstance(data, SE3):
+            scale = torch.ones_like(data.data[..., :1])
+            data = torch.cat([data.data, scale], -1)
+
+        elif isinstance(data, Sim3):
+            data = data.data
+
+        super(Sim3, self).__init__(data)
+
+
+def cat(group_list, dim):
+    """Concatenate groups along dimension"""
+    data = torch.cat([X.data for X in group_list], dim=dim)
+    return group_list[0].__class__(data)
+
+
+def stack(group_list, dim):
+    """Concatenate groups along dimension"""
+    data = torch.stack([X.data for X in group_list], dim=dim)
+    return group_list[0].__class__(data)

third_party/vipe/vipe/ext/scatter.py

@@ -0,0 +1,205 @@
+# This file includes code originally from the Pytorch Scatter repository:
+# https://github.com/rusty1s/pytorch_scatter
+# Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+
+from typing import Optional, Tuple
+
+import torch
+
+from . import scatter_ext
+
+
+def broadcast(src: torch.Tensor, other: torch.Tensor, dim: int):
+    if dim < 0:
+        dim = other.dim() + dim
+    if src.dim() == 1:
+        for _ in range(0, dim):
+            src = src.unsqueeze(0)
+    for _ in range(src.dim(), other.dim()):
+        src = src.unsqueeze(-1)
+    src = src.expand(other.size())
+    return src
+
+
+def scatter_sum(
+    src: torch.Tensor,
+    index: torch.Tensor,
+    dim: int = -1,
+    out: Optional[torch.Tensor] = None,
+    dim_size: Optional[int] = None,
+) -> torch.Tensor:
+    index = broadcast(index, src, dim)
+    if out is None:
+        size = list(src.size())
+        if dim_size is not None:
+            size[dim] = dim_size
+        elif index.numel() == 0:
+            size[dim] = 0
+        else:
+            size[dim] = int(index.max()) + 1
+        out = torch.zeros(size, dtype=src.dtype, device=src.device)
+        return out.scatter_add_(dim, index, src)
+    else:
+        return out.scatter_add_(dim, index, src)
+
+
+def scatter_add(
+    src: torch.Tensor,
+    index: torch.Tensor,
+    dim: int = -1,
+    out: Optional[torch.Tensor] = None,
+    dim_size: Optional[int] = None,
+) -> torch.Tensor:
+    return scatter_sum(src, index, dim, out, dim_size)
+
+
+def scatter_mul(
+    src: torch.Tensor,
+    index: torch.Tensor,
+    dim: int = -1,
+    out: Optional[torch.Tensor] = None,
+    dim_size: Optional[int] = None,
+) -> torch.Tensor:
+    return scatter_ext.scatter_mul(src, index, dim, out, dim_size)
+
+
+def scatter_mean(
+    src: torch.Tensor,
+    index: torch.Tensor,
+    dim: int = -1,
+    out: Optional[torch.Tensor] = None,
+    dim_size: Optional[int] = None,
+) -> torch.Tensor:
+    out = scatter_sum(src, index, dim, out, dim_size)
+    dim_size = out.size(dim)
+
+    index_dim = dim
+    if index_dim < 0:
+        index_dim = index_dim + src.dim()
+    if index.dim() <= index_dim:
+        index_dim = index.dim() - 1
+
+    ones = torch.ones(index.size(), dtype=src.dtype, device=src.device)
+    count = scatter_sum(ones, index, index_dim, None, dim_size)
+    count[count < 1] = 1
+    count = broadcast(count, out, dim)
+    if out.is_floating_point():
+        out.true_divide_(count)
+    else:
+        out.div_(count, rounding_mode="floor")
+    return out
+
+
+def scatter_min(
+    src: torch.Tensor,
+    index: torch.Tensor,
+    dim: int = -1,
+    out: Optional[torch.Tensor] = None,
+    dim_size: Optional[int] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    return scatter_ext.scatter_min(src, index, dim, out, dim_size)
+
+
+def scatter_max(
+    src: torch.Tensor,
+    index: torch.Tensor,
+    dim: int = -1,
+    out: Optional[torch.Tensor] = None,
+    dim_size: Optional[int] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    return scatter_ext.scatter_max(src, index, dim, out, dim_size)
+
+
+def scatter(
+    src: torch.Tensor,
+    index: torch.Tensor,
+    dim: int = -1,
+    out: Optional[torch.Tensor] = None,
+    dim_size: Optional[int] = None,
+    reduce: str = "sum",
+) -> torch.Tensor:
+    r"""
+    |
+
+    .. image:: https://raw.githubusercontent.com/rusty1s/pytorch_scatter/
+            master/docs/source/_figures/add.svg?sanitize=true
+        :align: center
+        :width: 400px
+
+    |
+
+    Reduces all values from the :attr:`src` tensor into :attr:`out` at the
+    indices specified in the :attr:`index` tensor along a given axis
+    :attr:`dim`.
+    For each value in :attr:`src`, its output index is specified by its index
+    in :attr:`src` for dimensions outside of :attr:`dim` and by the
+    corresponding value in :attr:`index` for dimension :attr:`dim`.
+    The applied reduction is defined via the :attr:`reduce` argument.
+
+    Formally, if :attr:`src` and :attr:`index` are :math:`n`-dimensional
+    tensors with size :math:`(x_0, ..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})`
+    and :attr:`dim` = `i`, then :attr:`out` must be an :math:`n`-dimensional
+    tensor with size :math:`(x_0, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})`.
+    Moreover, the values of :attr:`index` must be between :math:`0` and
+    :math:`y - 1`, although no specific ordering of indices is required.
+    The :attr:`index` tensor supports broadcasting in case its dimensions do
+    not match with :attr:`src`.
+
+    For one-dimensional tensors with :obj:`reduce="sum"`, the operation
+    computes
+
+    .. math::
+        \mathrm{out}_i = \mathrm{out}_i + \sum_j~\mathrm{src}_j
+
+    where :math:`\sum_j` is over :math:`j` such that
+    :math:`\mathrm{index}_j = i`.
+
+    .. note::
+
+        This operation is implemented via atomic operations on the GPU and is
+        therefore **non-deterministic** since the order of parallel operations
+        to the same value is undetermined.
+        For floating-point variables, this results in a source of variance in
+        the result.
+
+    :param src: The source tensor.
+    :param index: The indices of elements to scatter.
+    :param dim: The axis along which to index. (default: :obj:`-1`)
+    :param out: The destination tensor.
+    :param dim_size: If :attr:`out` is not given, automatically create output
+        with size :attr:`dim_size` at dimension :attr:`dim`.
+        If :attr:`dim_size` is not given, a minimal sized output tensor
+        according to :obj:`index.max() + 1` is returned.
+    :param reduce: The reduce operation (:obj:`"sum"`, :obj:`"mul"`,
+        :obj:`"mean"`, :obj:`"min"` or :obj:`"max"`). (default: :obj:`"sum"`)
+
+    :rtype: :class:`Tensor`
+
+    .. code-block:: python
+
+        from torch_scatter import scatter
+
+        src = torch.randn(10, 6, 64)
+        index = torch.tensor([0, 1, 0, 1, 2, 1])
+
+        # Broadcasting in the first and last dim.
+        out = scatter(src, index, dim=1, reduce="sum")
+
+        print(out.size())
+
+    .. code-block::
+
+        torch.Size([10, 3, 64])
+    """
+    if reduce == "sum" or reduce == "add":
+        return scatter_sum(src, index, dim, out, dim_size)
+    if reduce == "mul":
+        return scatter_mul(src, index, dim, out, dim_size)
+    elif reduce == "mean":
+        return scatter_mean(src, index, dim, out, dim_size)
+    elif reduce == "min":
+        return scatter_min(src, index, dim, out, dim_size)[0]
+    elif reduce == "max":
+        return scatter_max(src, index, dim, out, dim_size)[0]
+    else:
+        raise ValueError

third_party/vipe/vipe/ext/specs.py

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from pathlib import Path
+
+
+def get_sources() -> list[str]:
+    csrc_path = Path(__file__).parent.parent.parent / "csrc"
+    return [str(p) for p in csrc_path.glob("**/*") if p.suffix in [".cpp", ".cu"]]
+
+
+def _additional_include_flags() -> list[str]:
+    if "CONDA_PREFIX" in os.environ:
+        conda_include_path = os.path.join(os.environ["CONDA_PREFIX"], "include")
+        if os.path.exists(conda_include_path):
+            return ["-isystem", conda_include_path]
+    return []
+
+
+def get_cpp_flags() -> list[str]:
+    return ["-O3", "-DWITH_CUDA"] + _additional_include_flags()
+
+
+def get_cuda_flags() -> list[str]:
+    return ["-O3", "-DWITH_CUDA", "--use_fast_math"] + _additional_include_flags()

third_party/vipe/vipe/ext/xformers.py

@@ -0,0 +1,49 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Sequence
+
+import torch
+import torch.nn.functional as F
+
+
+def memory_efficient_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    # (B, S, H, D) -> (B, H, S, D)
+    query = query.transpose(1, 2)
+    key = key.transpose(1, 2)
+    value = value.transpose(1, 2)
+
+    out = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_bias)
+
+    return out.transpose(1, 2)
+
+
+def scaled_index_add(
+    input: torch.Tensor,  # [B, M, D]
+    index: torch.Tensor,  # [Bi] - int64
+    source: torch.Tensor,  # [Bi, M, D]
+    scaling: Optional[torch.Tensor] = None,  # [D]
+    alpha: float = 1.0,
+) -> torch.Tensor:
+    return torch.index_add(input, dim=0, source=scaling * source, index=index, alpha=alpha)
+
+
+def index_select_cat(sources: Sequence[torch.Tensor], indices: Sequence[torch.Tensor]) -> torch.Tensor:
+    return torch.cat([s[i.long()].flatten() for s, i in zip(sources, indices)], dim=0)

third_party/vipe/vipe/pipeline/__init__.py

@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import copy
+import importlib
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Sequence
+
+from omegaconf import DictConfig
+
+from vipe.streams.base import MultiviewVideoList, VideoStream
+
+
+@dataclass(kw_only=True, slots=True)
+class AnnotationPipelineOutput:
+    # Eager return of the payload values that comes from early stages of the pipeline.
+    payload: Any | None = None
+    output_streams: Sequence[VideoStream] | None = None
+
+
+class Pipeline(ABC):
+    def __init__(self) -> None:
+        self._return_payload = False
+        self._return_output_streams = False
+
+    @property
+    def return_payload(self) -> bool:
+        return self._return_payload
+
+    @return_payload.setter
+    def return_payload(self, value: bool) -> None:
+        assert isinstance(value, bool), "return_payload must be a boolean"
+        self._return_payload = value
+        if value:
+            self._return_output_streams = False
+
+    @property
+    def return_output_streams(self) -> bool:
+        return self._return_output_streams
+
+    @return_output_streams.setter
+    def return_output_streams(self, value: bool) -> None:
+        assert isinstance(value, bool), "return_output_streams must be a boolean"
+        self._return_output_streams = value
+        if value:
+            self._return_payload = False
+
+    def should_filter(self, stream_name: str) -> bool:
+        return False
+
+    @abstractmethod
+    def run(self, video_data: VideoStream | MultiviewVideoList) -> AnnotationPipelineOutput: ...
+
+
+def make_pipeline_cls(config: DictConfig) -> type[Pipeline]:
+    module_path, class_name = config.instance.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)
+
+
+def make_pipeline(config: DictConfig) -> Pipeline:
+    config = copy.deepcopy(config)
+    pipeline_cls = make_pipeline_cls(config)
+    del config.instance
+    return pipeline_cls(**config)

third_party/vipe/vipe/pipeline/default.py

@@ -0,0 +1,164 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import pickle
+
+from pathlib import Path
+
+import torch
+
+from omegaconf import DictConfig
+
+from vipe.slam.system import SLAMOutput, SLAMSystem
+from vipe.streams.base import (
+    AssignAttributesProcessor,
+    FrameAttribute,
+    MultiviewVideoList,
+    ProcessedVideoStream,
+    StreamProcessor,
+    VideoStream,
+)
+from vipe.utils import io
+from vipe.utils.cameras import CameraType
+from vipe.utils.visualization import save_projection_video
+
+from . import AnnotationPipelineOutput, Pipeline
+from .processors import (
+    AdaptiveDepthProcessor,
+    GeoCalibIntrinsicsProcessor,
+    MultiviewDepthProcessor,
+    TrackAnythingProcessor,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class DefaultAnnotationPipeline(Pipeline):
+    def __init__(self, init: DictConfig, slam: DictConfig, post: DictConfig, output: DictConfig) -> None:
+        super().__init__()
+        self.init_cfg = init
+        self.slam_cfg = slam
+        self.post_cfg = post
+        self.out_cfg = output
+        self.out_path = Path(self.out_cfg.path)
+        self.out_path.mkdir(exist_ok=True, parents=True)
+        self.camera_type = CameraType(self.init_cfg.camera_type)
+
+    def _add_init_processors(self, video_stream: VideoStream) -> ProcessedVideoStream:
+        init_processors: list[StreamProcessor] = []
+
+        # The assertions make sure that the attributes are not estimated previously.
+        # Otherwise it will be overwritten by the processors.
+        assert FrameAttribute.INTRINSICS not in video_stream.attributes()
+        assert FrameAttribute.CAMERA_TYPE not in video_stream.attributes()
+        assert FrameAttribute.METRIC_DEPTH not in video_stream.attributes()
+        assert FrameAttribute.INSTANCE not in video_stream.attributes()
+
+        init_processors.append(GeoCalibIntrinsicsProcessor(video_stream, camera_type=self.camera_type))
+        if self.init_cfg.instance is not None:
+            init_processors.append(
+                TrackAnythingProcessor(
+                    self.init_cfg.instance.phrases,
+                    add_sky=self.init_cfg.instance.add_sky,
+                    sam_run_gap=int(video_stream.fps() * self.init_cfg.instance.kf_gap_sec),
+                )
+            )
+        return ProcessedVideoStream(video_stream, init_processors)
+
+    def _add_post_processors(
+        self, view_idx: int, video_stream: VideoStream, slam_output: SLAMOutput
+    ) -> ProcessedVideoStream:
+        post_processors: list[StreamProcessor] = [
+            AssignAttributesProcessor(
+                {
+                    FrameAttribute.POSE: slam_output.get_view_trajectory(view_idx),  # type: ignore
+                    FrameAttribute.INTRINSICS: [slam_output.intrinsics[view_idx]] * len(video_stream),
+                }
+            )
+        ]
+        if (depth_align_model := self.post_cfg.depth_align_model) is not None:
+            if depth_align_model.startswith("mvd_"):
+                post_processors.append(MultiviewDepthProcessor(slam_output, model=depth_align_model))
+            else:
+                post_processors.append(AdaptiveDepthProcessor(slam_output, view_idx, depth_align_model))
+        return ProcessedVideoStream(video_stream, post_processors)
+
+    def run(self, video_data: VideoStream | MultiviewVideoList) -> AnnotationPipelineOutput:
+        if isinstance(video_data, MultiviewVideoList):
+            video_streams = [video_data[view_idx] for view_idx in range(len(video_data))]
+            artifact_paths = [io.ArtifactPath(self.out_path, video_stream.name()) for video_stream in video_streams]
+            slam_rig = video_data.rig()
+
+        else:
+            assert isinstance(video_data, VideoStream)
+            video_streams = [video_data]
+            artifact_paths = [io.ArtifactPath(self.out_path, video_data.name())]
+            slam_rig = None
+
+        annotate_output = AnnotationPipelineOutput()
+
+        if all([self.should_filter(video_stream.name()) for video_stream in video_streams]):
+            logger.info(f"{video_data.name()} has been proccessed already, skip it!!")
+            return annotate_output
+
+        slam_streams: list[VideoStream] = [
+            self._add_init_processors(video_stream).cache("process", online=True) for video_stream in video_streams
+        ]
+
+        slam_pipeline = SLAMSystem(device=torch.device("cuda"), config=self.slam_cfg)
+        slam_output = slam_pipeline.run(slam_streams, rig=slam_rig, camera_type=self.camera_type)
+
+        if self.return_payload:
+            annotate_output.payload = slam_output
+            return annotate_output
+
+        output_streams = [
+            self._add_post_processors(view_idx, slam_stream, slam_output).cache("depth", online=True)
+            for view_idx, slam_stream in enumerate(slam_streams)
+        ]
+
+        # Dumping artifacts for all views in the streams
+        for output_stream, artifact_path in zip(output_streams, artifact_paths):
+            artifact_path.meta_info_path.parent.mkdir(exist_ok=True, parents=True)
+            if self.out_cfg.save_artifacts:
+                logger.info(f"Saving artifacts to {artifact_path}")
+                if self.out_cfg.get("cameras_only", False):
+                    io.save_pose_artifacts(artifact_path, output_stream)
+                    io.save_intrinsics_artifacts(artifact_path, output_stream)
+                else:
+                    io.save_artifacts(artifact_path, output_stream)
+                with artifact_path.meta_info_path.open("wb") as f:
+                    pickle.dump({"ba_residual": slam_output.ba_residual}, f)
+
+            if self.out_cfg.save_viz:
+                save_projection_video(
+                    artifact_path.meta_vis_path,
+                    output_stream,
+                    slam_output,
+                    self.out_cfg.viz_downsample,
+                    self.out_cfg.viz_attributes,
+                )
+
+            if self.out_cfg.save_slam_map and slam_output.slam_map is not None:
+                logger.info(f"Saving SLAM map to {artifact_path.slam_map_path}")
+                slam_output.slam_map.save(artifact_path.slam_map_path)
+
+        if self.return_output_streams:
+            annotate_output.output_streams = output_streams
+
+        return annotate_output

third_party/vipe/vipe/pipeline/processors.py

@@ -0,0 +1,441 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+from typing import Iterator
+
+import numpy as np
+import torch
+
+from vipe.priors.depth import DepthEstimationInput, make_depth_model
+from vipe.priors.depth.alignment import align_inv_depth_to_depth
+from vipe.priors.depth.priorda import PriorDAModel
+from vipe.priors.depth.videodepthanything import VideoDepthAnythingDepthModel
+from vipe.priors.geocalib import GeoCalib
+from vipe.priors.track_anything import TrackAnythingPipeline
+from vipe.slam.interface import SLAMOutput
+from vipe.streams.base import (CachedVideoStream, FrameAttribute,
+                               StreamProcessor, VideoFrame, VideoStream)
+from vipe.utils.cameras import CameraType
+from vipe.utils.logging import pbar
+from vipe.utils.misc import unpack_optional
+from vipe.utils.morph import erode
+
+logger = logging.getLogger(__name__)
+
+
+class IntrinsicEstimationProcessor(StreamProcessor):
+    """Override existing intrinsics with estimated intrinsics."""
+
+    def __init__(self, video_stream: VideoStream, gap_sec: float = 1.0) -> None:
+        super().__init__()
+        gap_frame = int(gap_sec * video_stream.fps())
+        gap_frame = min(gap_frame, (len(video_stream) - 1) // 2)
+        self.sample_frame_inds = [0, gap_frame, gap_frame * 2]
+        self.fov_y = -1.0
+        self.camera_type = CameraType.PINHOLE
+        self.distortion: list[float] = []
+
+    def update_attributes(self, previous_attributes: set[FrameAttribute]) -> set[FrameAttribute]:
+        return previous_attributes | {FrameAttribute.INTRINSICS}
+
+    def __call__(self, frame_idx: int, frame: VideoFrame) -> VideoFrame:
+        assert self.fov_y > 0, "FOV not set"
+        frame_height, frame_width = frame.size()
+        fx = fy = frame_height / (2 * np.tan(self.fov_y / 2))
+        frame.intrinsics = torch.as_tensor(
+            [fx, fy, frame_width / 2, frame_height / 2] + self.distortion,
+        ).float()
+        frame.camera_type = self.camera_type
+        return frame
+
+
+class GeoCalibIntrinsicsProcessor(IntrinsicEstimationProcessor):
+    def __init__(
+        self,
+        video_stream: VideoStream,
+        gap_sec: float = 1.0,
+        camera_type: CameraType = CameraType.PINHOLE,
+    ) -> None:
+        super().__init__(video_stream, gap_sec)
+
+        is_pinhole = camera_type == CameraType.PINHOLE
+        weights = "pinhole" if is_pinhole else "distorted"
+
+        model = GeoCalib(weights=weights).cuda()
+        indexable_stream = CachedVideoStream(video_stream)
+
+        if is_pinhole:
+            sample_frames = torch.stack([indexable_stream[i].rgb.moveaxis(-1, 0) for i in self.sample_frame_inds])
+            res = model.calibrate(
+                sample_frames,
+                shared_intrinsics=True,
+            )
+        else:
+            # Use first frame for calibration
+            camera_model = {
+                CameraType.PINHOLE: "pinhole",
+                CameraType.MEI: "simple_mei",
+            }[camera_type]
+            res = model.calibrate(
+                indexable_stream[self.sample_frame_inds[0]].rgb.moveaxis(-1, 0)[None],
+                camera_model=camera_model,
+            )
+
+        self.fov_y = res["camera"].vfov[0].item()
+        self.camera_type = camera_type
+
+        if not is_pinhole:
+            # Assign distortion parameter
+            self.distortion = [res["camera"].dist[0, 0].item()]
+
+
+class TrackAnythingProcessor(StreamProcessor):
+    """
+    A processor that tracks a mask caption in the video.
+    """
+
+    def __init__(
+        self,
+        mask_phrases: list[str],
+        add_sky: bool,
+        sam_run_gap: int = 30,
+        mask_expand: int = 5,
+    ) -> None:
+        self.mask_phrases = mask_phrases
+        self.sam_run_gap = sam_run_gap
+        self.add_sky = add_sky
+
+        if self.add_sky:
+            self.mask_phrases.append(VideoFrame.SKY_PROMPT)
+
+        self.tracker = TrackAnythingPipeline(self.mask_phrases, sam_points_per_side=50, sam_run_gap=self.sam_run_gap)
+        self.mask_expand = mask_expand
+
+    def update_attributes(self, previous_attributes: set[FrameAttribute]) -> set[FrameAttribute]:
+        return previous_attributes | {FrameAttribute.INSTANCE, FrameAttribute.MASK}
+
+    def __call__(self, frame_idx: int, frame: VideoFrame) -> VideoFrame:
+        frame.instance, frame.instance_phrases = self.tracker.track(frame)
+        self.last_track_frame = frame.raw_frame_idx
+
+        frame_instance_mask = frame.instance == 0
+        if self.add_sky:
+            # We won't mask out the sky.
+            frame_instance_mask |= frame.sky_mask
+
+        frame.mask = erode(frame_instance_mask, self.mask_expand)
+        return frame
+
+
+class AdaptiveDepthProcessor(StreamProcessor):
+    """
+    Compute projection of the SLAM map onto the current frames.
+    If it's well-distributed, then use the fast map-prompted video depth model.
+    If not, then use the slow metric depth + video depth alignment model.
+    """
+
+    def __init__(
+        self,
+        slam_output: SLAMOutput,
+        view_idx: int = 0,
+        model: str = "adaptive_unidepth-l_svda",
+        share_depth_model: bool = False,
+    ):
+        super().__init__()
+        self.slam_output = slam_output
+        self.infill_target_pose = self.slam_output.get_view_trajectory(view_idx)
+        assert view_idx == 0, "Adaptive depth processor only supports view_idx=0"
+        assert not share_depth_model, "Adaptive depth processor does not support shared depth model"
+        self.require_cache = True
+        self.model = model
+
+        try:
+            prefix, metric_model, video_model = model.split("_")
+            assert video_model in ["svda", "vda"]
+            self.video_depth_model = VideoDepthAnythingDepthModel(model="vits" if video_model == "svda" else "vitl")
+
+        except ValueError:
+            prefix, metric_model = model.split("_")
+            video_model = None
+            self.video_depth_model = None
+
+        assert prefix == "adaptive", "Model name should start with 'adaptive_'"
+
+        self.depth_model = make_depth_model(metric_model)
+        self.prompt_model = PriorDAModel()
+        self.update_momentum = 0.99
+
+    def __call__(self, frame_idx: int, frame: VideoFrame) -> VideoFrame:
+        raise NotImplementedError("AdaptiveDepthProcessor should not be called directly.")
+
+    def update_attributes(self, previous_attributes: set[FrameAttribute]) -> set[FrameAttribute]:
+        return previous_attributes | {FrameAttribute.METRIC_DEPTH}
+
+    def _compute_uv_score(self, depth: torch.Tensor, patch_count: int = 10) -> float:
+        h_shape = depth.size(0) // patch_count
+        w_shape = depth.size(1) // patch_count
+        depth_crop = (depth > 0)[: h_shape * patch_count, : w_shape * patch_count]
+        depth_crop = depth_crop.reshape(patch_count, h_shape, patch_count, w_shape)
+        depth_exist = depth_crop.any(dim=(1, 3))
+        return depth_exist.float().mean().item()
+
+    def _compute_video_da(self, frame_iterator: Iterator[VideoFrame]) -> tuple[torch.Tensor, list[VideoFrame]]:
+        frame_list: list[np.ndarray] = []
+        frame_data_list: list[VideoFrame] = []
+        for frame in frame_iterator:
+            frame_data_list.append(frame.cpu())
+            frame_list.append(frame.rgb.cpu().numpy())
+
+        video_depth_result: torch.Tensor = unpack_optional(
+            self.video_depth_model.estimate(DepthEstimationInput(video_frame_list=frame_list)).relative_inv_depth
+        )
+        return video_depth_result, frame_data_list
+
+    def update_iterator(self, previous_iterator: Iterator[VideoFrame], pass_idx: int) -> Iterator[VideoFrame]:
+        # Determine the percentage score of the SLAM map.
+
+        self.cache_scale_bias = None
+        min_uv_score: float = 1.0
+
+        if self.video_depth_model is not None:
+            video_depth_result, data_iterator = self._compute_video_da(previous_iterator)
+        else:
+            video_depth_result = None
+            data_iterator = previous_iterator
+
+        for frame_idx, frame in pbar(enumerate(data_iterator), desc="Aligning depth"):
+            # Convert back to GPU if not already.
+            frame = frame.cuda()
+
+            # Compute the minimum UV score only once at the 0-th frame.
+            if frame_idx == 0:
+                for test_frame_idx in range(self.slam_output.trajectory.shape[0]):
+                    if test_frame_idx % 10 != 0:
+                        continue
+                    depth_infilled = self.slam_output.slam_map.project_map(
+                        test_frame_idx,
+                        0,
+                        frame.size(),
+                        unpack_optional(frame.intrinsics),
+                        self.infill_target_pose[test_frame_idx],
+                        unpack_optional(frame.camera_type),
+                        infill=False,
+                    )
+                    uv_score = self._compute_uv_score(depth_infilled)
+                    if uv_score < min_uv_score:
+                        min_uv_score = uv_score
+
+                logger.info(f"Minimum UV score: {min_uv_score:.4f}")
+
+            if min_uv_score < 0.3:
+                prompt_result = self.depth_model.estimate(
+                    DepthEstimationInput(
+                        rgb=frame.rgb.float().cuda(), intrinsics=frame.intrinsics, camera_type=frame.camera_type
+                    )
+                ).metric_depth
+                frame.information = f"uv={min_uv_score:.2f}(Metric)"
+            else:
+                depth_map = self.slam_output.slam_map.project_map(
+                    frame_idx,
+                    0,
+                    frame.size(),
+                    unpack_optional(frame.intrinsics),
+                    self.infill_target_pose[frame_idx],
+                    unpack_optional(frame.camera_type),
+                    infill=False,
+                )
+                if frame.mask is not None:
+                    depth_map = depth_map * frame.mask.float()
+                prompt_result = self.prompt_model.estimate(
+                    DepthEstimationInput(
+                        rgb=frame.rgb.float().cuda(),
+                        prompt_metric_depth=depth_map,
+                    )
+                ).metric_depth
+                frame.information = f"uv={min_uv_score:.2f}(SLAM)"
+
+            if video_depth_result is not None:
+                video_depth_inv_depth = video_depth_result[frame_idx]
+
+                align_mask = video_depth_inv_depth > 1e-3
+                if frame.mask is not None:
+                    align_mask = align_mask & frame.mask & (~frame.sky_mask)
+
+                try:
+                    _, scale, bias = align_inv_depth_to_depth(
+                        unpack_optional(video_depth_inv_depth),
+                        prompt_result,
+                        align_mask,
+                    )
+                except RuntimeError:
+                    scale, bias = self.cache_scale_bias
+
+                # momentum update
+                if self.cache_scale_bias is None:
+                    self.cache_scale_bias = (scale, bias)
+                scale = self.cache_scale_bias[0] * self.update_momentum + scale * (1 - self.update_momentum)
+                bias = self.cache_scale_bias[1] * self.update_momentum + bias * (1 - self.update_momentum)
+                self.cache_scale_bias = (scale, bias)
+
+                video_inv_depth = video_depth_inv_depth * scale + bias
+                video_inv_depth[video_inv_depth < 1e-3] = 1e-3
+                frame.metric_depth = video_inv_depth.reciprocal()
+
+            else:
+                frame.metric_depth = prompt_result
+
+            yield frame
+
+
+class MultiviewDepthProcessor(StreamProcessor):
+    """
+    Use multi-view depth model (e.g. DAv3, MapAnything, CAPA) to estimate depth map for each frame.
+    To ensure that the depth maps are consistent with the SLAM map/pose (metric), we condition the depth model either with
+    (a) sparse points, or (b) camera poses & intrinsics.
+
+    Depth is estimated in a sliding-window manner, and overlapped frames are linearly averaged to sharp transitions.
+    To create enough parallex to improve estimation confidence, for each window we optionally also include
+    neighboring keyframes, and their secondary neighboring keyframes.
+    (Multi-view input video frames are currently not supported)
+    """
+
+    def __init__(
+        self,
+        slam_output: SLAMOutput,
+        model: str = "mvd_dav3",
+        window_size: int = 10,                  # Practically this should be as large as possible if memory permits.
+        overlap_size: int = 3,
+        secondary_keyframe: bool = False,       # This is found to cause jittering for some scenes due to abrupt context changes.
+    ):
+        super().__init__()
+        self.slam_output = slam_output
+        self.model = model
+        self.window_size = window_size
+        self.overlap_size = overlap_size
+        self.secondary_keyframe = secondary_keyframe
+
+        self.keyframes_inds = unpack_optional(self.slam_output.slam_map).dense_disp_frame_inds
+        self.keyframes_data: list[VideoFrame] = []
+        self.n_frames = 0
+
+        # Need two passes for this iterator to work.
+        self.n_passes_required = 2
+
+        if self.model == "mvd_dav3":
+            try:
+                from depth_anything_3.api import DepthAnything3
+                from depth_anything_3.api import logger as dav3_logger
+            except ModuleNotFoundError:
+                raise ModuleNotFoundError(
+                    "depth-anything-3 not found. Please reinstall vipe with `pip install --no-build-isolation -e .[dav3]`"
+                )
+
+            dav3_logger.level = 0  # Disable logging timing information
+            self.dav3_api = DepthAnything3.from_pretrained("depth-anything/DA3-GIANT")
+            self.dav3_api = self.dav3_api.cuda().eval()
+
+    def update_attributes(self, previous_attributes: set[FrameAttribute]) -> set[FrameAttribute]:
+        return previous_attributes | {FrameAttribute.METRIC_DEPTH}
+
+    def __call__(self, frame_idx: int, frame: VideoFrame) -> VideoFrame:
+        raise NotImplementedError("MultiviewDepthProcessor should not be called directly.")
+
+    def _probe_keyframe_indices(self, frame_idx: int) -> list[int]:
+        inds: list[int] = []
+        left_idx = np.searchsorted(self.keyframes_inds, frame_idx, side="right").item() - 1
+        inds.append(left_idx)
+        if frame_idx < self.keyframes_inds[-1]:
+            inds.append(left_idx + 1)
+        # Pick the farthest secondary keyframe from the left keyframe.
+        if self.secondary_keyframe:
+            slam_graph = unpack_optional(self.slam_output.slam_map).backend_graph
+            if slam_graph is not None:
+                matching_secondary_j = slam_graph[slam_graph[:, 0] == left_idx, 1].tolist()
+                picked_sj_idx = np.argmax([abs(self.keyframes_inds[j] - frame_idx) for j in matching_secondary_j])
+                inds.append(matching_secondary_j[picked_sj_idx])
+        return inds
+
+    def record_keyframes(self, previous_iterator: Iterator[VideoFrame]) -> Iterator[VideoFrame]:
+        for frame_idx, frame in enumerate(previous_iterator):
+            self.n_frames += 1
+            if frame_idx in self.keyframes_inds:
+                self.keyframes_data.append(frame)
+            yield frame
+
+    def estimate_depth_sliding_window(self, previous_iterator: Iterator[VideoFrame]) -> Iterator[VideoFrame]:
+        current_sliding_window: list[VideoFrame] = []
+        current_sliding_window_idx: list[int] = []
+        trailing_depth: torch.Tensor | None = None
+        for frame_idx, frame in pbar(enumerate(previous_iterator), desc="Estimating multi-view depth"):
+            current_sliding_window.append(frame)
+            current_sliding_window_idx.append(frame_idx)
+            is_last_frame = frame_idx == self.n_frames - 1
+
+            if len(current_sliding_window) == self.window_size or is_last_frame:
+                # Grab all neighboring keyframes to anchor the current sliding window.
+                # Note that we remove redundant keyframes that already exist in the current sliding window.
+                sw_keyframe_inds = list(
+                    set(sum([self._probe_keyframe_indices(i) for i in current_sliding_window_idx], []))
+                )
+                sw_keyframe_inds = [
+                    t for t in sw_keyframe_inds if self.keyframes_inds[t] not in current_sliding_window_idx
+                ]
+
+                sw_images, sw_exts, sw_ints = zip(*[frame.dav3_conditions() for frame in current_sliding_window])
+
+                if len(sw_keyframe_inds) > 0:
+                    kf_images, kf_exts, kf_ints = zip(*[self.keyframes_data[t].dav3_conditions() for t in sw_keyframe_inds])
+                else:
+                    kf_images, kf_exts, kf_ints = tuple(), tuple(), tuple()
+
+                # Perform inference
+                dav3_inference_result = self.dav3_api.inference(
+                    list(sw_images + kf_images),
+                    extrinsics=np.stack(sw_exts + kf_exts, axis=0),
+                    intrinsics=np.stack(sw_ints + kf_ints, axis=0),
+                    process_res_method="lower_bound_resize",  # Keep aspect ratio
+                )
+                sw_depth = torch.from_numpy(dav3_inference_result.depth[: len(sw_images)]).float().cuda()
+                sw_depth = torch.nn.functional.interpolate(sw_depth[:, None], frame.size(), mode="bilinear")[:, 0]
+
+                n_frames_to_yield = (
+                    self.window_size - self.overlap_size if not is_last_frame else len(current_sliding_window)
+                )
+
+                # Linearly interpolate the trailing depth with new depth
+                if trailing_depth is not None:
+                    n_interp_frames = len(trailing_depth)
+                    alpha = torch.linspace(0, 1, n_interp_frames + 2)[1:-1].float().cuda()[:, None, None]
+                    sw_depth[:n_interp_frames] = trailing_depth * (1 - alpha) + sw_depth[:n_interp_frames] * alpha
+
+                for sw_idx, frame in enumerate(current_sliding_window[:n_frames_to_yield]):
+                    frame.metric_depth = sw_depth[sw_idx]
+                    yield frame
+
+                trailing_depth = sw_depth[n_frames_to_yield:]
+                current_sliding_window = current_sliding_window[n_frames_to_yield:]
+                current_sliding_window_idx = current_sliding_window_idx[n_frames_to_yield:]
+
+        assert len(current_sliding_window) == 0, "Current sliding window should be empty"
+
+    def update_iterator(self, previous_iterator: Iterator[VideoFrame], pass_idx: int) -> Iterator[VideoFrame]:
+        if pass_idx == 0:
+            yield from self.record_keyframes(previous_iterator)
+        elif pass_idx == 1:
+            yield from self.estimate_depth_sliding_window(previous_iterator)
+        else:
+            raise ValueError(f"Invalid pass index: {pass_idx}")

third_party/vipe/vipe/priors/depth/__init__.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .base import DepthEstimationInput, DepthEstimationModel, DepthEstimationResult, DepthType
+
+
+def make_depth_model(model: str):
+    if "-" not in model:
+        model_name, model_sub = model, ""
+    else:
+        model_name, model_sub = model.split("-")
+
+    if model_name == "metric3d":
+        from .metric3d import Metric3DDepthModel
+
+        return Metric3DDepthModel(version=2, model=model_sub)
+
+    elif model_name == "unidepth":
+        from .unidepth import UniDepth2Model
+
+        return UniDepth2Model(type=model_sub)
+
+    elif model_name == "moge":
+        from .moge import MogeModel
+
+        return MogeModel()
+
+    elif model_name == "dav3":
+        from .dav3 import DepthAnything3Model
+
+        return DepthAnything3Model()
+
+    else:
+        raise ValueError(f"Unknown depth model: {model}")

third_party/vipe/vipe/priors/depth/adapter.py

@@ -0,0 +1,105 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from vipe.utils.misc import unpack_optional
+from vipe.utils.cameras import CameraType, BaseCameraModel
+
+from .base import DepthEstimationInput, DepthEstimationModel, DepthEstimationResult, DepthType
+
+
+class PinholeDepthAdapter(DepthEstimationModel):
+    """
+    Adapter for depth estimation models that support only pinhole cameras.
+    Works by rectifying the input image to pinhole, perform depth estimation, and then unrectify the depth map.
+    """
+
+    def __init__(self, model: DepthEstimationModel):
+        super().__init__()
+        self.model = model
+        assert CameraType.PINHOLE in self.model.supported_camera_types, "Model must support pinhole cameras"
+
+        self.last_intrinsics: torch.Tensor | None = None
+        self.last_pinhole_intrinsics: torch.Tensor | None = None
+        self.last_rectification_map: tuple[torch.Tensor, torch.Tensor] | None = None
+
+    @property
+    def supported_camera_types(self) -> list[CameraType]:
+        return [CameraType.PINHOLE, CameraType.MEI]
+
+    @property
+    def depth_type(self) -> DepthType:
+        # Model needs to be re-run if intrinsics are updated
+        return DepthType.MODEL_METRIC_DEPTH
+
+    def _compute_rectification_map(self, h: int, w: int, src_model: BaseCameraModel, dst_model: BaseCameraModel):
+        # Compute rectification map that should be performed on the src img to obtain the dst model img
+        device = src_model.intrinsics.device
+        y, x = torch.meshgrid(torch.arange(h, device=device).float(), torch.arange(w, device=device).float(), indexing="ij")
+        pts, _, _ = dst_model.iproj_disp(torch.ones_like(x), x, y)
+        coords, _, _ = src_model.proj_points(pts)
+        coords_norm = 2.0 * coords / torch.tensor([w, h], device=coords.device) - 1.0
+        coords_norm = coords_norm.reshape(1, h, w, 2)
+        return coords_norm
+
+    def estimate(self, src: DepthEstimationInput) -> DepthEstimationResult:
+        # Short-cut for supported camera types
+        if src.camera_type in self.model.supported_camera_types:
+            return self.model.estimate(src)
+
+        rgb: torch.Tensor = unpack_optional(src.rgb)
+        if rgb.dim() == 3:
+            rgb, batch_dim = rgb[None], False
+        else:
+            batch_dim = True
+        img_h, img_w = rgb.shape[1:3]
+
+        if self.last_intrinsics is None or not torch.allclose(self.last_intrinsics, src.intrinsics):
+            tgt_intr = src.camera_type.build_camera_model(src.intrinsics)
+            pinhole_intr = tgt_intr.pinhole()
+            to_pinhole_map = self._compute_rectification_map(img_h, img_w, tgt_intr, pinhole_intr)
+            to_tgt_map = self._compute_rectification_map(img_h, img_w, pinhole_intr, tgt_intr)
+            self.last_intrinsics = src.intrinsics.clone()
+            self.last_pinhole_intrinsics = pinhole_intr.intrinsics.clone()
+            self.last_rectification_map = (to_pinhole_map, to_tgt_map)
+
+        assert self.last_rectification_map is not None
+        to_pinhole_map, to_tgt_map = self.last_rectification_map
+
+        rgb = rgb.permute(0, 3, 1, 2)
+        rgb = torch.nn.functional.grid_sample(
+            rgb,
+            to_pinhole_map,
+            mode="bilinear",
+            align_corners=False,
+        )
+        rgb = rgb.permute(0, 2, 3, 1)
+
+        est = self.model.estimate(
+            DepthEstimationInput(rgb=rgb, intrinsics=self.last_pinhole_intrinsics, camera_type=CameraType.PINHOLE)
+        )
+        metric_depth = torch.nn.functional.grid_sample(
+            est.metric_depth.unsqueeze(1),
+            to_tgt_map,
+            mode="nearest",
+            align_corners=False,
+        )
+        metric_depth = metric_depth.squeeze(1)
+
+        if not batch_dim:
+            metric_depth = metric_depth[0]
+
+        return DepthEstimationResult(metric_depth=metric_depth)

third_party/vipe/vipe/priors/depth/alignment.py

@@ -0,0 +1,125 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+
+def align_inv_depth_to_depth(
+    source_inv_depth: torch.Tensor,
+    target_depth: torch.Tensor,
+    target_mask: torch.Tensor | None = None,
+    quantile_masking: bool = True,
+) -> tuple[torch.Tensor, float, float]:
+    """
+    Apply affine transformation to align source inverse depth to target depth.
+
+    Args:
+        source_inv_depth: Inverse depth map to be aligned. Shape: (H, W).
+        target_depth: Target depth map. Shape: (H, W).
+        target_mask: Mask of valid target pixels. Shape: (H, W).
+
+    Returns:
+        Aligned Depth map. Shape: (H, W).
+        scale: Scaling factor.
+        bias: Bias term.
+    """
+    target_inv_depth = 1.0 / target_depth
+    source_mask = source_inv_depth > 0
+    target_depth_mask = target_depth > 0
+
+    if target_mask is None:
+        target_mask = target_depth_mask
+    else:
+        target_mask = torch.logical_and(target_mask > 0, target_depth_mask)
+
+    # Remove outliers
+    if quantile_masking:
+        outlier_quantiles = torch.tensor([0.1, 0.9], device=source_inv_depth.device)
+        source_data_low, source_data_high = torch.quantile(source_inv_depth[source_mask], outlier_quantiles)
+        target_data_low, target_data_high = torch.quantile(target_inv_depth[target_mask], outlier_quantiles)
+        source_mask = (source_inv_depth > source_data_low) & (source_inv_depth < source_data_high)
+        target_mask = (target_inv_depth > target_data_low) & (target_inv_depth < target_data_high)
+
+    mask = torch.logical_and(source_mask, target_mask)
+
+    source_data = source_inv_depth[mask].view(-1, 1)
+    target_data = target_inv_depth[mask].view(-1, 1)
+
+    ones = torch.ones((source_data.shape[0], 1), device=source_data.device)
+    source_data_h = torch.cat([source_data, ones], dim=1)
+    transform_matrix = torch.linalg.lstsq(source_data_h, target_data).solution
+
+    scale, bias = transform_matrix[0, 0], transform_matrix[1, 0]
+    aligned_inv_depth = source_inv_depth * scale + bias
+    aligned_depth = torch.clamp(aligned_inv_depth.reciprocal(), min=1e-4)
+
+    return aligned_depth, scale, bias
+
+
+def align_depth_to_depth(
+    source_depth: torch.Tensor,
+    target_depth: torch.Tensor,
+    target_mask: torch.Tensor | None = None,
+    quantile_masking: bool = True,
+    bias: bool = True,
+) -> torch.Tensor:
+    """
+    Apply affine transformation to align source depth to target depth.
+
+    Args:
+        source_inv_depth: Depth map to be aligned. Shape: (H, W).
+        target_depth: Target depth map. Shape: (H, W).
+        target_mask: Mask of valid target pixels. Shape: (H, W).
+
+    Returns:
+        Aligned Depth map. Shape: (H, W).
+    """
+    source_mask = source_depth > 0
+    target_depth_mask = target_depth > 0
+
+    if target_mask is None:
+        target_mask = target_depth_mask
+    else:
+        target_mask = torch.logical_and(target_mask > 0, target_depth_mask)
+
+    # Remove outliers
+    if quantile_masking:
+        outlier_quantiles = torch.tensor([0.1, 0.9], device=source_depth.device)
+
+        source_data_low, source_data_high = torch.quantile(source_depth[source_mask], outlier_quantiles)
+        target_data_low, target_data_high = torch.quantile(target_depth[target_mask], outlier_quantiles)
+        source_mask = (source_depth > source_data_low) & (source_depth < source_data_high)
+        target_mask = (target_depth > target_data_low) & (target_depth < target_data_high)
+
+    mask = torch.logical_and(source_mask, target_mask)
+
+    source_data = source_depth[mask].view(-1, 1)
+    target_data = target_depth[mask].view(-1, 1)
+
+    if not bias:
+        ones = torch.ones((source_data.shape[0], 1), device=source_data.device)
+        source_data_h = torch.cat([source_data, ones], dim=1)
+        transform_matrix = torch.linalg.lstsq(source_data_h, target_data).solution
+
+        scale, bias = transform_matrix[0, 0], transform_matrix[1, 0]
+        aligned_depth = source_depth * scale + bias
+        aligned_depth = torch.clamp(aligned_depth, min=1e-4)
+
+    else:
+        scale = torch.median(target_data / source_data)
+        aligned_depth = source_depth * scale
+        aligned_depth = torch.clamp(aligned_depth, min=1e-4)
+
+    return aligned_depth

third_party/vipe/vipe/priors/depth/base.py

@@ -0,0 +1,114 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+
+import numpy as np
+import torch
+
+from vipe.utils.cameras import CameraType
+
+
+class DepthType(Enum):
+    """
+    Type of depth estimated.
+    """
+
+    # Inverse metric depth (DepthPro, Metric3D), scale is determined by focal.
+    METRIC_DEPTH = "metric_depth"
+
+    # Metric depth (UniDepth), scale is determined by focal but require re-running the model.
+    MODEL_METRIC_DEPTH = "model_metric_depth"
+
+    # Metric distance, scale is determined by focal but require re-running the model.
+    MODEL_METRIC_DISTANCE = "model_metric_distance"
+
+    # Affine-invariant inverse depth (DepthAnything, MoGE), affine needs to be solved per-estimation.
+    AFFINE_DISP = "affine_disp"
+
+    # Scale-invariant depth (DUSt3R), scale needs to be solved per-estimation.
+    SCALE_DISP = "scale_disp"
+
+
+@dataclass(slots=True, kw_only=True)
+class DepthEstimationResult:
+    """
+    Dataclass for depth estimation results.
+
+    - relative_inv_depth: The estimated inverse depth map ([B,], H, W) in relative scale.
+        One has to estimate scale and offset of such estimation
+    - metric_depth: The estimated depth map ([B,], H, W) in metric scale.
+    - confidence: The confidence map ([B,], H, W).
+    """
+
+    relative_inv_depth: torch.Tensor | None = None
+    metric_depth: torch.Tensor | None = None
+    confidence: torch.Tensor | None = None
+
+
+@dataclass(slots=True, kw_only=True)
+class DepthEstimationInput:
+    """
+    Dataclass for depth estimation inputs.
+
+    - rgb: The source image ([B,], H, W, 3), should be within 0-1 float.
+    - video_frame_list: The list of video frames (H, W, 3), should be within 0-1 float, length is T.
+        we use numpy here mainly to enforce CPU tensor.
+    - intrinsics: The intrinsics of the camera.
+    - camera_type: The type of camera.
+    """
+
+    rgb: torch.Tensor | None = None
+    video_frame_list: list[np.ndarray] | None = None
+    prompt_metric_depth: torch.Tensor | None = None
+    intrinsics: torch.Tensor | None = None
+    camera_type: CameraType = CameraType.PINHOLE
+
+
+class DepthEstimationModel(ABC):
+    """
+    Unified interface for depth prediction models.
+    """
+
+    @property
+    def depth_type(self) -> DepthType:
+        """
+        Type of depth estimated.
+        """
+        raise NotImplementedError
+
+    @property
+    def supported_camera_types(self) -> list[CameraType]:
+        """
+        Supported camera types.
+        """
+        return [CameraType.PINHOLE]
+
+    @abstractmethod
+    def estimate(self, src: DepthEstimationInput) -> DepthEstimationResult:
+        """
+        Estimate a single optical flow result from two images.
+        """
+
+
+class DummyDepthModel(DepthEstimationModel):
+    """
+    Dummy model that would not predict anything.
+    """
+
+    def estimate(self, src: DepthEstimationInput) -> DepthEstimationResult:
+        return DepthEstimationResult()

third_party/vipe/vipe/priors/depth/dav2/__init__.py

@@ -0,0 +1,105 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+import torch
+
+from vipe.utils.misc import unpack_optional
+
+from ..base import DepthEstimationInput, DepthEstimationModel, DepthEstimationResult, DepthType
+from .dpt import DepthAnythingV2
+
+
+class DepthAnythingDepthModel(DepthEstimationModel):
+    """
+    https://github.com/DepthAnything/Depth-Anything-V2
+    """
+
+    def __init__(self, model: str = "vitl", ckpt: str = "default") -> None:
+        super().__init__()
+
+        self.model_config = {
+            "vits": {
+                "encoder": "vits",
+                "features": 64,
+                "out_channels": [48, 96, 192, 384],
+            },
+            "vitb": {
+                "encoder": "vitb",
+                "features": 128,
+                "out_channels": [96, 192, 384, 768],
+            },
+            "vitl": {
+                "encoder": "vitl",
+                "features": 256,
+                "out_channels": [256, 512, 1024, 1024],
+            },
+            "vitg": {
+                "encoder": "vitg",
+                "features": 384,
+                "out_channels": [1536, 1536, 1536, 1536],
+            },
+        }[model]
+        url_name = {"vits": "Small", "vitb": "Base", "vitl": "Large", "vitg": "Giant"}[model]
+
+        if ckpt == "default":
+            self.is_metric = False
+            self.max_depth = None
+            self.ckpt_url = f"https://huggingface.co/depth-anything/Depth-Anything-V2-{url_name}/resolve/main/depth_anything_v2_{model}.pth?download=true"
+
+        elif ckpt == "metric-indoor":
+            self.is_metric = True
+            self.max_depth = 20.0
+            self.ckpt_url = f"https://huggingface.co/depth-anything/Depth-Anything-V2-Metric-Hypersim-{url_name}/resolve/main/depth_anything_v2_metric_hypersim_{model}.pth?download=true"
+
+        elif ckpt == "metric-outdoor":
+            self.is_metric = True
+            self.max_depth = 80.0
+            self.ckpt_url = f"hhttps://huggingface.co/depth-anything/Depth-Anything-V2-Metric-VKITTI-{url_name}/resolve/main/depth_anything_v2_metric_vkitti_{model}.pth?download=true"
+
+        else:
+            raise ValueError("Invalid checkpoint name.")
+
+        self.model = DepthAnythingV2(**self.model_config, max_depth=self.max_depth)
+        self.model.load_state_dict(torch.hub.load_state_dict_from_url(self.ckpt_url, map_location="cpu"))
+        self.model.cuda().eval()
+
+    @property
+    def depth_type(self) -> DepthType:
+        return DepthType.AFFINE_DISP
+
+    def estimate(self, src: DepthEstimationInput) -> DepthEstimationResult:
+        rgb: torch.Tensor = unpack_optional(src.rgb)
+        assert rgb.dtype == torch.float32, "Input image should be float32"
+
+        if rgb.dim() == 3:
+            rgb, batch_dim = rgb[None], False
+        else:
+            batch_dim = True
+
+        assert rgb.size(0) == 1, "Only batch size 1 is supported."
+
+        # TODO: Support pytorch pre-processing
+        depth = self.model.infer_image(rgb[0].cpu().numpy())
+
+        if batch_dim:
+            depth = depth[None]
+
+        if self.is_metric:
+            return DepthEstimationResult(metric_depth=depth)
+
+        else:
+            return DepthEstimationResult(relative_inv_depth=depth)

third_party/vipe/vipe/priors/depth/dav2/dinov2.py

@@ -0,0 +1,429 @@
+# 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 logging
+import math
+
+from functools import partial
+from typing import Callable, Sequence, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+
+from torch.nn.init import trunc_normal_
+
+from .dinov2_layers import MemEffAttention, Mlp, PatchEmbed
+from .dinov2_layers import NestedTensorBlock as Block
+
+
+logger = logging.getLogger(__name__)
+
+
+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,
+            )
+            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
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        # DINOv2 with register modify the interpolate_offset from 0.1 to 0.0
+        w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
+        # w0, h0 = w0 + 0.1, h0 + 0.1
+
+        sqrt_N = math.sqrt(N)
+        sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
+            scale_factor=(sx, sy),
+            # (int(w0), int(h0)), # to solve the upsampling shape issue
+            mode="bicubic",
+            antialias=self.interpolate_antialias,
+        )
+
+        assert int(w0) == patch_pos_embed.shape[-2]
+        assert int(h0) == patch_pos_embed.shape[-1]
+        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:
+            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:
+            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
+
+
+def DINOv2(model_name):
+    model_zoo = {
+        "vits": vit_small,
+        "vitb": vit_base,
+        "vitl": vit_large,
+        "vitg": vit_giant2,
+    }
+
+    return model_zoo[model_name](
+        img_size=518,
+        patch_size=14,
+        init_values=1.0,
+        ffn_layer="mlp" if model_name != "vitg" else "swiglufused",
+        block_chunks=0,
+        num_register_tokens=0,
+        interpolate_antialias=False,
+        interpolate_offset=0.1,
+    )

third_party/vipe/vipe/priors/depth/dav2/dinov2_layers/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .attention import MemEffAttention
+from .block import NestedTensorBlock
+from .mlp import Mlp
+from .patch_embed import PatchEmbed
+from .swiglu_ffn import SwiGLUFFN

third_party/vipe/vipe/priors/depth/dav2/dinov2_layers/attention.py

@@ -0,0 +1,71 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import logging
+
+import torch
+
+from torch import Tensor, nn
+
+from vipe.ext.xformers import memory_efficient_attention
+
+
+logger = logging.getLogger(__name__)
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+        attn = q @ k.transpose(-2, -1)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(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)
+
+        q, k, v = torch.unbind(qkv, dim=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

third_party/vipe/vipe/priors/depth/dav2/dinov2_layers/block.py

@@ -0,0 +1,253 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+import logging
+
+from typing import Any, Callable, Dict, List, Tuple
+
+import torch
+
+from torch import Tensor, nn
+
+from vipe.ext.xformers import index_select_cat, scaled_index_add
+
+from .attention import Attention, MemEffAttention
+from .drop_path import DropPath
+from .layer_scale import LayerScale
+from .mlp import Mlp
+
+
+# from xformers.ops import fmha
+
+
+logger = logging.getLogger(__name__)
+
+
+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):
+            assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
+            return self.forward_nested(x_or_x_list)
+        else:
+            raise AssertionError

third_party/vipe/vipe/priors/depth/dav2/dinov2_layers/drop_path.py

@@ -0,0 +1,37 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
+
+
+from torch import nn
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)

third_party/vipe/vipe/priors/depth/dav2/dinov2_layers/layer_scale.py

@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
+
+from typing import Union
+
+import torch
+from torch import Tensor, nn
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.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

third_party/vipe/vipe/priors/depth/dav2/dinov2_layers/mlp.py

@@ -0,0 +1,41 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
+
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x

third_party/vipe/vipe/priors/depth/dav2/dinov2_layers/patch_embed.py

@@ -0,0 +1,114 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+from typing import Callable, Optional, Tuple, Union
+
+import torch.nn as nn
+from torch import Tensor
+
+
+def make_2tuple(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class PatchEmbed(nn.Module):
+    """
+    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
+
+    Args:
+        img_size: Image size.
+        patch_size: Patch token size.
+        in_chans: Number of input image channels.
+        embed_dim: Number of linear projection output channels.
+        norm_layer: Normalization layer.
+    """
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int]] = 224,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        super().__init__()
+
+        image_HW = make_2tuple(img_size)
+        patch_HW = make_2tuple(patch_size)
+        patch_grid_size = (
+            image_HW[0] // patch_HW[0],
+            image_HW[1] // patch_HW[1],
+        )
+
+        self.img_size = image_HW
+        self.patch_size = patch_HW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW
+        )
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def init_alpha_conv(self, cond_channels=2):
+        self.alpha_proj = nn.Conv2d(
+            cond_channels,
+            self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+        )
+        nn.init.constant_(self.alpha_proj.weight, 0)
+        nn.init.constant_(self.alpha_proj.bias, 0)
+
+    def forward(self, x: Tensor, condition=None) -> 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
+        if hasattr(self, "alpha_proj"):
+            x += self.alpha_proj(condition)
+
+        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

third_party/vipe/vipe/priors/depth/dav2/dinov2_layers/swiglu_ffn.py

@@ -0,0 +1,34 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Callable, Optional
+
+import torch.nn.functional as F
+
+from torch import Tensor, nn
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)

third_party/vipe/vipe/priors/depth/dav2/dpt.py

@@ -0,0 +1,281 @@
+# This file includes code originally from the Depth Anything V2 repository:
+# https://github.com/DepthAnything/Depth-Anything-V2
+# Licensed under the Apache License, Version 2.0. See THIRD_PARTY_LICENSES.md for details.
+
+
+import cv2
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torchvision.transforms import Compose
+
+from .dinov2 import DINOv2
+from .util.blocks import FeatureFusionBlock, _make_scratch
+from .util.transform import NormalizeImage, PrepareForNet, Resize
+
+
+def _make_fusion_block(features, use_bn, size=None):
+    return FeatureFusionBlock(
+        features,
+        nn.ReLU(False),
+        deconv=False,
+        bn=use_bn,
+        expand=False,
+        align_corners=True,
+        size=size,
+    )
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_feature, out_feature):
+        super().__init__()
+
+        self.conv_block = nn.Sequential(
+            nn.Conv2d(in_feature, out_feature, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_feature),
+            nn.ReLU(True),
+        )
+
+    def forward(self, x):
+        return self.conv_block(x)
+
+
+class DPTHead(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        features=256,
+        use_bn=False,
+        out_channels=[256, 512, 1024, 1024],
+        use_clstoken=False,
+        is_metric=False,
+    ):
+        super(DPTHead, self).__init__()
+
+        self.use_clstoken = use_clstoken
+
+        self.projects = nn.ModuleList(
+            [
+                nn.Conv2d(
+                    in_channels=in_channels,
+                    out_channels=out_channel,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                )
+                for out_channel in out_channels
+            ]
+        )
+
+        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,
+                ),
+            ]
+        )
+
+        if use_clstoken:
+            self.readout_projects = nn.ModuleList()
+            for _ in range(len(self.projects)):
+                self.readout_projects.append(nn.Sequential(nn.Linear(2 * in_channels, in_channels), nn.GELU()))
+
+        self.scratch = _make_scratch(
+            out_channels,
+            features,
+            groups=1,
+            expand=False,
+        )
+
+        self.scratch.stem_transpose = None
+
+        self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
+
+        head_features_1 = features
+        head_features_2 = 32
+
+        self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)
+
+        if is_metric:
+            self.scratch.output_conv2 = nn.Sequential(
+                nn.Conv2d(
+                    head_features_1 // 2,
+                    head_features_2,
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                ),
+                nn.ReLU(True),
+                nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
+                nn.Sigmoid(),
+            )
+
+        else:
+            self.scratch.output_conv2 = nn.Sequential(
+                nn.Conv2d(
+                    head_features_1 // 2,
+                    head_features_2,
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                ),
+                nn.ReLU(True),
+                nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
+                nn.ReLU(True),
+                nn.Identity(),
+            )
+
+    def forward(self, out_features, patch_h, patch_w):
+        out = []
+        for i, x in enumerate(out_features):
+            if self.use_clstoken:
+                x, cls_token = x[0], x[1]
+                readout = cls_token.unsqueeze(1).expand_as(x)
+                x = self.readout_projects[i](torch.cat((x, readout), -1))
+            else:
+                x = x[0]
+
+            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
+
+            x = self.projects[i](x)
+            x = self.resize_layers[i](x)
+
+            out.append(x)
+
+        layer_1, layer_2, layer_3, layer_4 = out
+
+        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)
+
+        path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
+        path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
+        path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
+        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
+
+        out = self.scratch.output_conv1(path_1)
+        out = F.interpolate(
+            out,
+            (int(patch_h * 14), int(patch_w * 14)),
+            mode="bilinear",
+            align_corners=True,
+        )
+        out = self.scratch.output_conv2(out)
+
+        return out
+
+
+class DepthAnythingV2(nn.Module):
+    def __init__(
+        self,
+        encoder="vitl",
+        features=256,
+        out_channels=[256, 512, 1024, 1024],
+        use_bn=False,
+        use_clstoken=False,
+        max_depth: float | None = 20.0,
+    ):
+        super(DepthAnythingV2, self).__init__()
+
+        self.intermediate_layer_idx = {
+            "vits": [2, 5, 8, 11],
+            "vitb": [2, 5, 8, 11],
+            "vitl": [4, 11, 17, 23],
+            "vitg": [9, 19, 29, 39],
+        }
+
+        self.max_depth = max_depth
+
+        self.encoder = encoder
+        self.pretrained = DINOv2(model_name=encoder)
+
+        self.depth_head = DPTHead(
+            self.pretrained.embed_dim,
+            features,
+            use_bn,
+            out_channels=out_channels,
+            use_clstoken=use_clstoken,
+            is_metric=self.max_depth is not None,
+        )
+
+    def forward(self, x):
+        patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14
+
+        features = self.pretrained.get_intermediate_layers(
+            x, self.intermediate_layer_idx[self.encoder], return_class_token=True
+        )
+
+        depth = self.depth_head(features, patch_h, patch_w)
+
+        if self.max_depth is not None:
+            depth = self.max_depth * depth
+
+        else:
+            depth = F.relu(depth)
+
+        return depth.squeeze(1)
+
+    @torch.no_grad()
+    def infer_image(self, raw_image, input_size=518):
+        image, (h, w) = self.image2tensor(raw_image, input_size)
+
+        depth = self.forward(image)
+
+        depth = F.interpolate(depth[:, None], (h, w), mode="bilinear", align_corners=True)[0, 0]
+
+        # return depth.cpu().numpy()
+        return depth
+
+    def image2tensor(self, raw_image, input_size=518):
+        transform = Compose(
+            [
+                Resize(
+                    width=input_size,
+                    height=input_size,
+                    resize_target=False,
+                    keep_aspect_ratio=True,
+                    ensure_multiple_of=14,
+                    resize_method="lower_bound",
+                    image_interpolation_method=cv2.INTER_CUBIC,
+                ),
+                NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                PrepareForNet(),
+            ]
+        )
+
+        h, w = raw_image.shape[:2]
+
+        # image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0
+        image = raw_image
+
+        image = transform({"image": image})["image"]
+        image = torch.from_numpy(image).unsqueeze(0)
+
+        DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+        image = image.to(DEVICE)
+
+        return image, (h, w)

third_party/vipe/vipe/priors/depth/dav2/util/__init__.py

@@ -0,0 +1,3 @@
+# This file includes code originally from the Depth Anything V2 repository:
+# https://github.com/DepthAnything/Depth-Anything-V2
+# Licensed under the Apache License, Version 2.0. See THIRD_PARTY_LICENSES.md for details.

third_party/vipe/vipe/priors/depth/dav2/util/blocks.py

@@ -0,0 +1,207 @@
+# This file includes code originally from the Depth Anything V2 repository:
+# https://github.com/DepthAnything/Depth-Anything-V2
+# Licensed under the Apache License, Version 2.0. See THIRD_PARTY_LICENSES.md for details.
+
+
+import torch.nn as nn
+
+
+def _make_scratch(in_shape, out_shape, groups=1, expand=False):
+    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):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super().__init__()
+
+        self.bn = bn
+
+        self.groups = 1
+
+        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,
+        )
+
+        if self.bn == True:
+            self.bn1 = nn.BatchNorm2d(features)
+            self.bn2 = nn.BatchNorm2d(features)
+
+        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.bn == True:
+            out = self.bn1(out)
+
+        out = self.activation(out)
+        out = self.conv2(out)
+        if self.bn == True:
+            out = self.bn2(out)
+
+        if self.groups > 1:
+            out = self.conv_merge(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,
+    ):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super(FeatureFusionBlock, self).__init__()
+
+        self.deconv = deconv
+        self.align_corners = align_corners
+
+        self.groups = 1
+
+        self.expand = expand
+        out_features = features
+        if self.expand == True:
+            out_features = features // 2
+
+        self.out_conv = nn.Conv2d(
+            features,
+            out_features,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+            groups=1,
+        )
+
+        self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
+        self.resConfUnit2 = ResidualConvUnit(features, activation, bn)
+
+        self.skip_add = nn.quantized.FloatFunctional()
+
+        self.size = size
+
+    def forward(self, *xs, size=None):
+        """Forward pass.
+
+        Returns:
+            tensor: output
+        """
+        output = xs[0]
+
+        if len(xs) == 2:
+            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 = nn.functional.interpolate(output, **modifier, mode="bilinear", align_corners=self.align_corners)
+
+        output = self.out_conv(output)
+
+        return output

third_party/vipe/vipe/priors/depth/dav2/util/transform.py

@@ -0,0 +1,167 @@
+# This file includes code originally from the Depth Anything V2 repository:
+# https://github.com/DepthAnything/Depth-Anything-V2
+# Licensed under the Apache License, Version 2.0. See THIRD_PARTY_LICENSES.md for details.
+
+import cv2
+import numpy as np
+
+
+class Resize(object):
+    """Resize sample to given size (width, height)."""
+
+    def __init__(
+        self,
+        width,
+        height,
+        resize_target=True,
+        keep_aspect_ratio=False,
+        ensure_multiple_of=1,
+        resize_method="lower_bound",
+        image_interpolation_method=cv2.INTER_AREA,
+    ):
+        """Init.
+
+        Args:
+            width (int): desired output width
+            height (int): desired output height
+            resize_target (bool, optional):
+                True: Resize the full sample (image, mask, target).
+                False: Resize image only.
+                Defaults to True.
+            keep_aspect_ratio (bool, optional):
+                True: Keep the aspect ratio of the input sample.
+                Output sample might not have the given width and height, and
+                resize behaviour depends on the parameter 'resize_method'.
+                Defaults to False.
+            ensure_multiple_of (int, optional):
+                Output width and height is constrained to be multiple of this parameter.
+                Defaults to 1.
+            resize_method (str, optional):
+                "lower_bound": Output will be at least as large as the given size.
+                "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
+                "minimal": Scale as least as possible.  (Output size might be smaller than given size.)
+                Defaults to "lower_bound".
+        """
+        self.__width = width
+        self.__height = height
+
+        self.__resize_target = resize_target
+        self.__keep_aspect_ratio = keep_aspect_ratio
+        self.__multiple_of = ensure_multiple_of
+        self.__resize_method = resize_method
+        self.__image_interpolation_method = image_interpolation_method
+
+    def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
+        y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        if max_val is not None and y > max_val:
+            y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        if y < min_val:
+            y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        return y
+
+    def get_size(self, width, height):
+        # determine new height and width
+        scale_height = self.__height / height
+        scale_width = self.__width / width
+
+        if self.__keep_aspect_ratio:
+            if self.__resize_method == "lower_bound":
+                # scale such that output size is lower bound
+                if scale_width > scale_height:
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            elif self.__resize_method == "upper_bound":
+                # scale such that output size is upper bound
+                if scale_width < scale_height:
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            elif self.__resize_method == "minimal":
+                # scale as least as possbile
+                if abs(1 - scale_width) < abs(1 - scale_height):
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            else:
+                raise ValueError(f"resize_method {self.__resize_method} not implemented")
+
+        if self.__resize_method == "lower_bound":
+            new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
+            new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
+        elif self.__resize_method == "upper_bound":
+            new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
+            new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
+        elif self.__resize_method == "minimal":
+            new_height = self.constrain_to_multiple_of(scale_height * height)
+            new_width = self.constrain_to_multiple_of(scale_width * width)
+        else:
+            raise ValueError(f"resize_method {self.__resize_method} not implemented")
+
+        return (new_width, new_height)
+
+    def __call__(self, sample):
+        width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
+
+        # resize sample
+        sample["image"] = cv2.resize(
+            sample["image"],
+            (width, height),
+            interpolation=self.__image_interpolation_method,
+        )
+
+        if self.__resize_target:
+            if "depth" in sample:
+                sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
+
+            if "mask" in sample:
+                sample["mask"] = cv2.resize(
+                    sample["mask"].astype(np.float32),
+                    (width, height),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+
+        return sample
+
+
+class NormalizeImage(object):
+    """Normlize image by given mean and std."""
+
+    def __init__(self, mean, std):
+        self.__mean = mean
+        self.__std = std
+
+    def __call__(self, sample):
+        sample["image"] = (sample["image"] - self.__mean) / self.__std
+
+        return sample
+
+
+class PrepareForNet(object):
+    """Prepare sample for usage as network input."""
+
+    def __init__(self):
+        pass
+
+    def __call__(self, sample):
+        image = np.transpose(sample["image"], (2, 0, 1))
+        sample["image"] = np.ascontiguousarray(image).astype(np.float32)
+
+        if "depth" in sample:
+            depth = sample["depth"].astype(np.float32)
+            sample["depth"] = np.ascontiguousarray(depth)
+
+        if "mask" in sample:
+            sample["mask"] = sample["mask"].astype(np.float32)
+            sample["mask"] = np.ascontiguousarray(sample["mask"])
+
+        return sample

third_party/vipe/vipe/priors/depth/dav3.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+
+try:
+    from depth_anything_3.api import DepthAnything3
+    from depth_anything_3.api import logger as dav3_logger
+except ModuleNotFoundError:
+    DepthAnything3 = dav3_logger = None
+
+from vipe.utils.cameras import CameraType
+from vipe.utils.misc import unpack_optional
+
+from .base import DepthEstimationInput, DepthEstimationModel, DepthEstimationResult, DepthType
+
+
+class DepthAnything3Model(DepthEstimationModel):
+    """
+    https://github.com/ByteDance-Seed/Depth-Anything-3
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        if DepthAnything3 is None or dav3_logger is None:
+            raise RuntimeError(
+                "depth-anything-3 is not found. You can install it via `pip install --no-build-isolation -e .[dav3]`"
+            )
+
+        dav3_logger.level = 0  # Disable logging timing information
+        self.model = DepthAnything3.from_pretrained("depth-anything/DA3METRIC-LARGE")
+        self.model = self.model.cuda().eval()
+
+    @property
+    def depth_type(self) -> DepthType:
+        """
+        DAv3 model offers metric depth proportional to focal length
+        See: https://github.com/ByteDance-Seed/Depth-Anything-3?tab=readme-ov-file#-faq
+        """
+        return DepthType.METRIC_DEPTH
+
+    def estimate(self, src: DepthEstimationInput) -> DepthEstimationResult:
+        rgb: torch.Tensor = unpack_optional(src.rgb)
+        assert rgb.dtype == torch.float32, "Input image should be float32"
+
+        assert src.camera_type == CameraType.PINHOLE, "DAv3 only supports pinhole cameras"
+        focal_length: float = unpack_optional(src.intrinsics)[0].item()
+
+        if rgb.dim() == 3:
+            rgb, batch_dim = rgb[None], False
+        else:
+            batch_dim = True
+
+        rgb_images = [(rgb[idx].cpu().numpy() * 255).astype(np.uint8) for idx in range(rgb.shape[0])]
+
+        with torch.no_grad():
+            dav3_inference_result = self.model.inference(
+                rgb_images,
+                process_res_method="upper_bound_resize",
+                process_res=504,
+            )
+
+        # Compute focal internally in DAv3 model
+        dav3_camera_focal = focal_length / max(rgb.shape[2], rgb.shape[1]) * 504
+        # Normalize to correct scale, 300 is the magic/commonly used number chosen by the model authors.
+        dav3_metric_depth = dav3_inference_result.depth * dav3_camera_focal / 300.0
+        dav3_sky_mask = dav3_inference_result.sky
+
+        # Mark sky region as invalid depth, so that inaccurate predictions are not used for further computation.
+        dav3_metric_depth = dav3_metric_depth * (~dav3_sky_mask).astype(dav3_metric_depth.dtype)
+        dav3_metric_depth = torch.from_numpy(dav3_metric_depth).cuda()[None]
+        dav3_metric_depth = torch.nn.functional.interpolate(dav3_metric_depth, rgb.shape[1:3], mode="nearest")[:, 0]
+
+        if not batch_dim:
+            dav3_metric_depth = dav3_metric_depth.squeeze(0)
+
+        return DepthEstimationResult(metric_depth=dav3_metric_depth)

third_party/vipe/vipe/priors/depth/metric3d/__init__.py

@@ -0,0 +1,167 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+import torch
+
+from vipe.utils.misc import unpack_optional
+from vipe.utils.cameras import CameraType
+
+from ..base import DepthEstimationInput, DepthEstimationModel, DepthEstimationResult, DepthType
+from .model_fn import (
+    metric3d_convnext_large,
+    metric3d_convnext_tiny,
+    metric3d_vit_giant2,
+    metric3d_vit_large,
+    metric3d_vit_small,
+)
+
+
+class Metric3DDepthModel(DepthEstimationModel):
+    """
+    https://github.com/YvanYin/Metric3D
+    """
+
+    @dataclass
+    class PadInfo:
+        pad_up: int
+        pad_down: int
+        pad_left: int
+        pad_right: int
+        size_origin: tuple[int, int]
+        focal_length: float
+
+    def __init__(self, version: int = 2, model: str = "giant2") -> None:
+        super().__init__()
+
+        if version == 1:
+            assert model in ["tiny", "large"]
+            if model == "tiny":
+                self.model = metric3d_convnext_tiny(pretrain=True, skip_validation=True)
+            else:
+                self.model = metric3d_convnext_large(pretrain=True, skip_validation=True)
+
+            self.input_size = (544, 1216)
+
+        elif version == 2:
+            assert model in ["small", "large", "giant2"]
+            if model == "small":
+                self.model = metric3d_vit_small(pretrain=True, skip_validation=True)
+            elif model == "large":
+                self.model = metric3d_vit_large(pretrain=True, skip_validation=True)
+            else:
+                self.model = metric3d_vit_giant2(pretrain=True, skip_validation=True)
+
+            self.input_size = (616, 1064)
+
+        else:
+            raise ValueError("Invalid version number.")
+
+        self.model.cuda().eval()
+
+    @property
+    def depth_type(self) -> DepthType:
+        return DepthType.METRIC_DEPTH
+
+    @torch.no_grad()
+    def _prepare_input(self, img: torch.Tensor, focal_length: float) -> tuple[torch.Tensor, PadInfo]:
+        """
+        Convert a 0-1 RGB image to normalized and padded image required by the model.
+        input: (B, H, W, 3) RGB image
+        return: (B, 3, H', W') tensor and pad_info
+        """
+
+        h_origin, w_origin = img.shape[1:3]
+        scale = min(self.input_size[0] / h_origin, self.input_size[1] / w_origin)
+        h_scaled, w_scaled = int(h_origin * scale), int(w_origin * scale)
+
+        img = img.permute(0, 3, 1, 2)
+        img = torch.nn.functional.interpolate(img, (h_scaled, w_scaled), mode="bilinear")
+
+        # normalize
+        mean = (
+            torch.tensor([123.675, 116.28, 103.53], device=img.device, dtype=torch.float32)[None, :, None, None] / 255
+        )
+        std = torch.tensor([58.395, 57.12, 57.375], device=img.device, dtype=torch.float32)[None, :, None, None] / 255
+        img = torch.div((img - mean), std)
+
+        # padding to input_size
+        pad_h = self.input_size[0] - h_scaled
+        pad_w = self.input_size[1] - w_scaled
+
+        pad_up, pad_left = pad_h // 2, pad_w // 2
+        pad_down, pad_right = pad_h - pad_up, pad_w - pad_left
+
+        img = torch.nn.functional.pad(
+            img,
+            (pad_left, pad_right, pad_up, pad_down),
+            value=0.0,
+        )
+
+        return img, self.PadInfo(
+            pad_up,
+            pad_down,
+            pad_left,
+            pad_right,
+            (h_origin, w_origin),
+            focal_length * scale,
+        )
+
+    def _post_process(self, feature: torch.Tensor, pad_info: PadInfo, is_depth: bool) -> torch.Tensor:
+        """
+        Post-process the feature map to remove padding, and de-canonicalize the depth.
+        input: (B, C, H', W') tensor
+        return: (B, C, H, W) tensor
+        """
+
+        feature = feature[
+            ...,
+            pad_info.pad_up : feature.size(-2) - pad_info.pad_down,
+            pad_info.pad_left : feature.size(-1) - pad_info.pad_right,
+        ]
+        feature = torch.nn.functional.interpolate(feature, pad_info.size_origin, mode="bilinear")
+
+        if is_depth:
+            canonical_to_real_scale = pad_info.focal_length / 1000.0
+            feature = feature * canonical_to_real_scale
+
+        return feature
+
+    def estimate(self, src: DepthEstimationInput) -> DepthEstimationResult:
+        rgb: torch.Tensor = unpack_optional(src.rgb)
+        assert rgb.dtype == torch.float32, "Input image should be float32"
+
+        assert src.camera_type == CameraType.PINHOLE, "Metric3D only supports pinhole cameras"
+        focal_length: float = unpack_optional(src.intrinsics)[0].item()
+
+        if rgb.dim() == 3:
+            rgb, batch_dim = rgb[None], False
+        else:
+            batch_dim = True
+
+        rgb, pad_info = self._prepare_input(rgb, focal_length)
+        pred_depth, confidence, output_dict = self.model.inference({"input": rgb})
+
+        pred_depth = self._post_process(pred_depth, pad_info, is_depth=True)
+        confidence = self._post_process(confidence, pad_info, is_depth=False)
+
+        if not batch_dim:
+            pred_depth, confidence = pred_depth[0], confidence[0]
+
+        return DepthEstimationResult(
+            metric_depth=pred_depth[0],
+            confidence=confidence[0],
+        )

third_party/vipe/vipe/priors/depth/metric3d/model/__init__.py

@@ -0,0 +1,3 @@
+# This file includes code originally from the Metric3D repository:
+# https://github.com/YvanYin/Metric3D
+# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.

third_party/vipe/vipe/priors/depth/metric3d/model/backbones/ConvNeXt.py

@@ -0,0 +1,310 @@
+# This file includes code originally from the Metric3D repository:
+# https://github.com/YvanYin/Metric3D
+# Licensed under the BSD-2 License. See THIRD_PARTY_LICENSES.md for details.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from timm.models.layers import DropPath, trunc_normal_
+from timm.models.registry import register_model
+
+
+class Block(nn.Module):
+    r"""ConvNeXt Block. There are two equivalent implementations:
+    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
+    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
+    We use (2) as we find it slightly faster in PyTorch
+
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+    """
+
+    def __init__(self, dim, drop_path=0.0, layer_scale_init_value=1e-6):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+        return x
+
+
+class ConvNeXt(nn.Module):
+    r"""ConvNeXt
+        A PyTorch impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+    Args:
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(
+        self,
+        in_chans=3,
+        num_classes=1000,
+        depths=[3, 3, 9, 3],
+        dims=[96, 192, 384, 768],
+        drop_path_rate=0.0,
+        layer_scale_init_value=1e-6,
+        head_init_scale=1.0,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.downsample_layers = nn.ModuleList()  # stem and 3 intermediate downsampling conv layers
+        stem = nn.Sequential(
+            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+            LayerNorm(dims[0], eps=1e-6, data_format="channels_first"),
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = nn.Sequential(
+                LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
+                nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = nn.ModuleList()  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(4):
+            stage = nn.Sequential(
+                *[
+                    Block(
+                        dim=dims[i],
+                        drop_path=dp_rates[cur + j],
+                        layer_scale_init_value=layer_scale_init_value,
+                    )
+                    for j in range(depths[i])
+                ]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        # self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer
+        # self.head = nn.Linear(dims[-1], num_classes)
+
+        self.apply(self._init_weights)
+        # self.head.weight.data.mul_(head_init_scale)
+        # self.head.bias.data.mul_(head_init_scale)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=0.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward_features(self, x):
+        features = []
+        for i in range(4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+            features.append(x)
+        return features  # global average pooling, (N, C, H, W) -> (N, C)
+
+    def forward(self, x):
+        # x = self.forward_features(x)
+        # x = self.head(x)
+        features = self.forward_features(x)
+        return features
+
+
+class LayerNorm(nn.Module):
+    r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
+    with shape (batch_size, channels, height, width).
+    """
+
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+
+model_urls = {
+    "convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
+    "convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
+    "convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
+    "convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
+    "convnext_tiny_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth",
+    "convnext_small_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth",
+    "convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
+    "convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
+    "convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
+}
+
+
+def convnext_tiny(pretrained=True, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
+    if pretrained:
+        checkpoint = torch.load(kwargs["checkpoint"], map_location="cpu")
+        # url = model_urls['convnext_tiny_22k'] if in_22k else model_urls['convnext_tiny_1k']
+        # checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
+        model_dict = model.state_dict()
+        pretrained_dict = {}
+        unmatched_pretrained_dict = {}
+        for k, v in checkpoint["model"].items():
+            if k in model_dict:
+                pretrained_dict[k] = v
+            else:
+                unmatched_pretrained_dict[k] = v
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+        print(
+            "Successfully loaded pretrained %d params, and %d paras are unmatched."
+            % (len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys()))
+        )
+        print("Unmatched pretrained paras are :", unmatched_pretrained_dict.keys())
+    return model
+
+
+def convnext_small(pretrained=True, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
+    if pretrained:
+        checkpoint = torch.load(kwargs["checkpoint"], map_location="cpu")
+        # url = model_urls['convnext_small_22k'] if in_22k else model_urls['convnext_small_1k']
+        # checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
+        model_dict = model.state_dict()
+        pretrained_dict = {}
+        unmatched_pretrained_dict = {}
+        for k, v in checkpoint["model"].items():
+            if k in model_dict:
+                pretrained_dict[k] = v
+            else:
+                unmatched_pretrained_dict[k] = v
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+        print(
+            "Successfully loaded pretrained %d params, and %d paras are unmatched."
+            % (len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys()))
+        )
+        print("Unmatched pretrained paras are :", unmatched_pretrained_dict.keys())
+    return model
+
+
+def convnext_base(pretrained=True, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
+    if pretrained:
+        checkpoint = torch.load(kwargs["checkpoint"], map_location="cpu")
+        # url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']
+        # checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
+        model_dict = model.state_dict()
+        pretrained_dict = {}
+        unmatched_pretrained_dict = {}
+        for k, v in checkpoint["model"].items():
+            if k in model_dict:
+                pretrained_dict[k] = v
+            else:
+                unmatched_pretrained_dict[k] = v
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+        print(
+            "Successfully loaded pretrained %d params, and %d paras are unmatched."
+            % (len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys()))
+        )
+        print("Unmatched pretrained paras are :", unmatched_pretrained_dict.keys())
+    return model
+
+
+def convnext_large(pretrained=True, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
+    if pretrained:
+        checkpoint = torch.load(kwargs["checkpoint"], map_location="cpu")
+        # url = model_urls['convnext_large_22k'] if in_22k else model_urls['convnext_large_1k']
+        # checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
+        model_dict = model.state_dict()
+        pretrained_dict = {}
+        unmatched_pretrained_dict = {}
+        for k, v in checkpoint["model"].items():
+            if k in model_dict:
+                pretrained_dict[k] = v
+            else:
+                unmatched_pretrained_dict[k] = v
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+        print(
+            "Successfully loaded pretrained %d params, and %d paras are unmatched."
+            % (len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys()))
+        )
+        print("Unmatched pretrained paras are :", unmatched_pretrained_dict.keys())
+    return model
+
+
+def convnext_xlarge(pretrained=True, in_22k=False, **kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
+    if pretrained:
+        assert in_22k, "only ImageNet-22K pre-trained ConvNeXt-XL is available; please set in_22k=True"
+        checkpoint = torch.load(kwargs["checkpoint"], map_location="cpu")
+        # url = model_urls['convnext_xlarge_22k']
+        # checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
+        model_dict = model.state_dict()
+        pretrained_dict = {}
+        unmatched_pretrained_dict = {}
+        for k, v in checkpoint["model"].items():
+            if k in model_dict:
+                pretrained_dict[k] = v
+            else:
+                unmatched_pretrained_dict[k] = v
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+        print(
+            "Successfully loaded pretrained %d params, and %d paras are unmatched."
+            % (len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys()))
+        )
+        print("Unmatched pretrained paras are :", unmatched_pretrained_dict.keys())
+    return model
+
+
+if __name__ == "__main__":
+    import torch
+
+    model = convnext_base(True, in_22k=False).cuda()
+
+    rgb = torch.rand((2, 3, 256, 256)).cuda()
+    out = model(rgb)
+    print(len(out))
+    for i, ft in enumerate(out):
+        print(i, ft.shape)