Merging main update from original repository

nerfstudio-project · Dec 14, 2024 · a3f628c · a3f628c
2 parents b0e978d + ff2da2e
commit a3f628c
Show file tree

Hide file tree

Showing 45 changed files with 930 additions and 883 deletions.
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
@@ -106,5 +106,6 @@ jobs:
       env:
         PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
       run: |         
+        BUILD_NO_CUDA=1 python -m build
         twine upload --username __token__ --password $PYPI_TOKEN dist/*
-      shell: bash
+      shell: bash
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -1 +1,2 @@
 recursive-include gsplat/cuda/csrc *
+recursive-include gsplat/cuda/include *
diff --git a/examples/simple_trainer_2dgs.py b/examples/simple_trainer_2dgs.py
@@ -436,7 +436,7 @@ def rasterize_splats(
                 **kwargs,
             )
         elif self.model_type == "2dgs-inria":
-            render_colors, render_alphas, info = rasterization_2dgs_inria_wrapper(
+            renders, info = rasterization_2dgs_inria_wrapper(
                 means=means,
                 quats=quats,
                 scales=scales,
@@ -577,6 +577,10 @@ def train(self):
                 step=step,
                 info=info,
             )
+            masks = data["mask"].to(device) if "mask" in data else None
+            if masks is not None:
+                pixels = pixels * masks[..., None]
+                colors = colors * masks[..., None]
 
             # loss
             l1loss = F.l1_loss(colors, pixels)

diff --git a/examples/simple_viewer.py b/examples/simple_viewer.py
@@ -68,7 +68,7 @@ def main(local_rank: int, world_rank, world_size: int, args):
                 quats,  # [N, 4]
                 scales,  # [N, 3]
                 opacities,  # [N]
-                colors,  # [N, 3]
+                colors,  # [N, S, 3]
                 viewmats,  # [C, 4, 4]
                 Ks,  # [C, 3, 3]
                 width,
@@ -181,7 +181,7 @@ def viewer_render_fn(camera_state: nerfview.CameraState, img_wh: Tuple[int, int]
             quats,  # [N, 4]
             scales,  # [N, 3]
             opacities,  # [N]
-            colors,  # [N, 3]
+            colors,  # [N, S, 3]
             viewmat[None],  # [1, 4, 4]
             K[None],  # [1, 3, 3]
             width,

diff --git a/gsplat/cuda/_backend.py b/gsplat/cuda/_backend.py
@@ -89,7 +89,7 @@ def cuda_toolkit_version():
         current_dir = os.path.dirname(os.path.abspath(__file__))
         glm_path = os.path.join(current_dir, "csrc", "third_party", "glm")
 
-        extra_include_paths = [os.path.join(PATH, "csrc/"), glm_path]
+        extra_include_paths = [os.path.join(PATH, "include/"), glm_path]
         extra_cflags = ["-O3"]
         if NO_FAST_MATH:
             extra_cuda_cflags = ["-O3"]

diff --git a/gsplat/cuda/csrc/adam.cu b/gsplat/cuda/csrc/adam.cu
@@ -1,6 +1,4 @@
 #include "bindings.h"
-#include "helpers.cuh"
-#include "utils.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/compute_sh_bwd.cu b/gsplat/cuda/csrc/compute_sh_bwd.cu
@@ -1,4 +1,5 @@
 #include "bindings.h"
+#include "helpers.cuh"
 #include "spherical_harmonics.cuh"
 #include "types.cuh"
 

diff --git a/gsplat/cuda/csrc/fully_fused_projection_2dgs_bwd.cu b/gsplat/cuda/csrc/fully_fused_projection_2dgs_bwd.cu
@@ -1,6 +1,7 @@
 #include "bindings.h"
 #include "helpers.cuh"
-#include "utils.cuh"
+#include "transform.cuh"
+#include "2dgs.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/fully_fused_projection_2dgs_fwd.cu b/gsplat/cuda/csrc/fully_fused_projection_2dgs_fwd.cu
@@ -1,6 +1,7 @@
 #include "bindings.h"
 #include "helpers.cuh"
-#include "utils.cuh"
+#include "transform.cuh"
+#include "2dgs.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/fully_fused_projection_bwd.cu b/gsplat/cuda/csrc/fully_fused_projection_bwd.cu
@@ -1,6 +1,10 @@
 #include "bindings.h"
 #include "helpers.cuh"
 #include "utils.cuh"
+#include "quat.cuh"
+#include "quat_scale_to_covar_preci.cuh"
+#include "proj.cuh"
+#include "transform.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/fully_fused_projection_fwd.cu b/gsplat/cuda/csrc/fully_fused_projection_fwd.cu
@@ -1,6 +1,8 @@
 #include "bindings.h"
-#include "helpers.cuh"
 #include "utils.cuh"
+#include "quat_scale_to_covar_preci.cuh"
+#include "proj.cuh"
+#include "transform.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/fully_fused_projection_packed_2dgs_bwd.cu b/gsplat/cuda/csrc/fully_fused_projection_packed_2dgs_bwd.cu
@@ -1,6 +1,7 @@
 #include "bindings.h"
 #include "helpers.cuh"
-#include "utils.cuh"
+#include "transform.cuh"
+#include "2dgs.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/fully_fused_projection_packed_2dgs_fwd.cu b/gsplat/cuda/csrc/fully_fused_projection_packed_2dgs_fwd.cu
@@ -1,7 +1,8 @@
 
 #include "bindings.h"
 #include "helpers.cuh"
-#include "utils.cuh"
+#include "quat.cuh"
+#include "transform.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/fully_fused_projection_packed_bwd.cu b/gsplat/cuda/csrc/fully_fused_projection_packed_bwd.cu
@@ -1,6 +1,10 @@
 #include "bindings.h"
 #include "helpers.cuh"
 #include "utils.cuh"
+#include "quat.cuh"
+#include "quat_scale_to_covar_preci.cuh"
+#include "proj.cuh"
+#include "transform.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/fully_fused_projection_packed_fwd.cu b/gsplat/cuda/csrc/fully_fused_projection_packed_fwd.cu
@@ -1,6 +1,8 @@
 #include "bindings.h"
-#include "helpers.cuh"
 #include "utils.cuh"
+#include "quat_scale_to_covar_preci.cuh"
+#include "proj.cuh"
+#include "transform.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/isect_tiles.cu b/gsplat/cuda/csrc/isect_tiles.cu
@@ -1,5 +1,4 @@
 #include "bindings.h"
-#include "helpers.cuh"
 #include "types.cuh"
 #include <cooperative_groups.h>
 #include <cub/cub.cuh>

diff --git a/gsplat/cuda/csrc/proj_bwd.cu b/gsplat/cuda/csrc/proj_bwd.cu
@@ -1,6 +1,5 @@
 #include "bindings.h"
-#include "helpers.cuh"
-#include "utils.cuh"
+#include "proj.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/proj_fwd.cu b/gsplat/cuda/csrc/proj_fwd.cu
@@ -1,6 +1,5 @@
 #include "bindings.h"
-#include "helpers.cuh"
-#include "utils.cuh"
+#include "proj.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/quat_scale_to_covar_preci_bwd.cu b/gsplat/cuda/csrc/quat_scale_to_covar_preci_bwd.cu
@@ -1,6 +1,6 @@
 #include "bindings.h"
-#include "helpers.cuh"
-#include "utils.cuh"
+#include "quat.cuh"
+#include "quat_scale_to_covar_preci.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/quat_scale_to_covar_preci_fwd.cu b/gsplat/cuda/csrc/quat_scale_to_covar_preci_fwd.cu
@@ -1,6 +1,5 @@
 #include "bindings.h"
-#include "helpers.cuh"
-#include "utils.cuh"
+#include "quat_scale_to_covar_preci.cuh"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>

diff --git a/gsplat/cuda/csrc/rasterize_to_indices_in_range.cu b/gsplat/cuda/csrc/rasterize_to_indices_in_range.cu
@@ -1,5 +1,4 @@
 #include "bindings.h"
-#include "helpers.cuh"
 #include "types.cuh"
 #include <cooperative_groups.h>
 #include <cub/cub.cuh>

diff --git a/gsplat/cuda/csrc/rasterize_to_indices_in_range_2dgs.cu b/gsplat/cuda/csrc/rasterize_to_indices_in_range_2dgs.cu
@@ -1,7 +1,6 @@
 #include "bindings.h"
-#include "helpers.cuh"
 #include "types.cuh"
-#include "utils.cuh"
+#include "2dgs.cuh"
 #include <cooperative_groups.h>
 #include <cub/cub.cuh>
 #include <cuda_runtime.h>

diff --git a/gsplat/cuda/csrc/rasterize_to_pixels_2dgs_bwd.cu b/gsplat/cuda/csrc/rasterize_to_pixels_2dgs_bwd.cu
@@ -1,7 +1,7 @@
 #include "bindings.h"
 #include "helpers.cuh"
 #include "types.cuh"
-#include "utils.cuh"
+#include "2dgs.cuh"
 #include <cooperative_groups.h>
 #include <cub/cub.cuh>
 #include <cuda_runtime.h>

diff --git a/gsplat/cuda/csrc/rasterize_to_pixels_2dgs_fwd.cu b/gsplat/cuda/csrc/rasterize_to_pixels_2dgs_fwd.cu
@@ -1,7 +1,6 @@
 #include "bindings.h"
-#include "helpers.cuh"
 #include "types.cuh"
-#include "utils.cuh"
+#include "2dgs.cuh"
 #include <cooperative_groups.h>
 #include <cub/cub.cuh>
 #include <cuda_runtime.h>

diff --git a/gsplat/cuda/csrc/rasterize_to_pixels_fwd.cu b/gsplat/cuda/csrc/rasterize_to_pixels_fwd.cu
@@ -1,5 +1,4 @@
 #include "bindings.h"
-#include "helpers.cuh"
 #include "types.cuh"
 #include <cooperative_groups.h>
 #include <cub/cub.cuh>

diff --git a/gsplat/cuda/csrc/third_party/glm b/gsplat/cuda/csrc/third_party/glm
Original file line number	Diff line number	Diff line change
		@@ -1 +1,2 @@
		recursive-include gsplat/cuda/csrc *
		recursive-include gsplat/cuda/include *