From 4e317c8d311b9220680036c0ed661c0861c3629a Mon Sep 17 00:00:00 2001 From: "Mr. For Example" <62230687+MrForExample@users.noreply.github.com> Date: Wed, 3 Jul 2024 17:52:40 -0600 Subject: [PATCH] Replace torchmcubes with PyMCubes torchmcubes dependencies is archived on github, so for better maintenance, I replace it with PyMCubes inside TripoSR --- Dockerfile | 2 +- Libs/simple-knn/ext.cpp | 17 ++ Libs/simple-knn/setup.py | 35 +++ Libs/simple-knn/simple_knn.cu | 222 ++++++++++++++++++++ Libs/simple-knn/simple_knn.h | 21 ++ Libs/simple-knn/simple_knn/.gitkeep | 0 Libs/simple-knn/spatial.cu | 26 +++ Libs/simple-knn/spatial.h | 14 ++ gen_3d_modules/InstantMesh/models/lrm.py | 4 +- gen_3d_modules/TripoSR/models/isosurface.py | 8 +- gen_3d_modules/TripoSR/system.py | 4 +- mesh_processer/mesh_utils.py | 14 +- requirements_post.txt | 3 +- requirements_post_win_py310_cu121.txt | 3 +- requirements_post_win_py311_cu121.txt | 3 +- requirements_post_win_py312_cu121.txt | 3 +- 16 files changed, 356 insertions(+), 23 deletions(-) create mode 100644 Libs/simple-knn/ext.cpp create mode 100644 Libs/simple-knn/setup.py create mode 100644 Libs/simple-knn/simple_knn.cu create mode 100644 Libs/simple-knn/simple_knn.h create mode 100644 Libs/simple-knn/simple_knn/.gitkeep create mode 100644 Libs/simple-knn/spatial.cu create mode 100644 Libs/simple-knn/spatial.h diff --git a/Dockerfile b/Dockerfile index f977674..7afb921 100644 --- a/Dockerfile +++ b/Dockerfile @@ -62,7 +62,7 @@ RUN pip install --no-cache -r requirements.txt WORKDIR /app/custom_nodes/ComfyUI-3D-Pack/ COPY --chown=user:user requirements.txt requirements_post.txt ./ -COPY --chown=user:user simple-knn/ simple-knn/ +COPY --chown=user:user Libs/simple-knn/ Libs/simple-knn/ COPY --chown=user:user Gen_3D_Modules/TriplaneGaussian/ Gen_3D_Modules/TriplaneGaussian/ RUN pip install --no-cache -r requirements.txt \ # post requirements installation require gpu, setup diff --git a/Libs/simple-knn/ext.cpp b/Libs/simple-knn/ext.cpp new file mode 100644 index 0000000..ae6cefe --- /dev/null +++ b/Libs/simple-knn/ext.cpp @@ -0,0 +1,17 @@ +/* + * Copyright (C) 2023, Inria + * GRAPHDECO research group, https://team.inria.fr/graphdeco + * All rights reserved. + * + * This software is free for non-commercial, research and evaluation use + * under the terms of the LICENSE.md file. + * + * For inquiries contact george.drettakis@inria.fr + */ + +#include +#include "spatial.h" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("distCUDA2", &distCUDA2); +} diff --git a/Libs/simple-knn/setup.py b/Libs/simple-knn/setup.py new file mode 100644 index 0000000..580d2bd --- /dev/null +++ b/Libs/simple-knn/setup.py @@ -0,0 +1,35 @@ +# +# Copyright (C) 2023, Inria +# GRAPHDECO research group, https://team.inria.fr/graphdeco +# All rights reserved. +# +# This software is free for non-commercial, research and evaluation use +# under the terms of the LICENSE.md file. +# +# For inquiries contact george.drettakis@inria.fr +# + +from setuptools import setup +from torch.utils.cpp_extension import CUDAExtension, BuildExtension +import os + +cxx_compiler_flags = [] + +if os.name == 'nt': + cxx_compiler_flags.append("/wd4624") + +setup( + name="simple_knn", + ext_modules=[ + CUDAExtension( + name="simple_knn._C", + sources=[ + "spatial.cu", + "simple_knn.cu", + "ext.cpp"], + extra_compile_args={"nvcc": [], "cxx": cxx_compiler_flags}) + ], + cmdclass={ + 'build_ext': BuildExtension + } +) diff --git a/Libs/simple-knn/simple_knn.cu b/Libs/simple-knn/simple_knn.cu new file mode 100644 index 0000000..880a49b --- /dev/null +++ b/Libs/simple-knn/simple_knn.cu @@ -0,0 +1,222 @@ +/* + * Copyright (C) 2023, Inria + * GRAPHDECO research group, https://team.inria.fr/graphdeco + * All rights reserved. + * + * This software is free for non-commercial, research and evaluation use + * under the terms of the LICENSE.md file. + * + * For inquiries contact george.drettakis@inria.fr + */ + +#define BOX_SIZE 1024 + +#include "cuda_runtime.h" +#include "device_launch_parameters.h" +#include "simple_knn.h" +#include +#include +#include +#include +#include +#include +#define __CUDACC__ +#include +#include +#include + +namespace cg = cooperative_groups; + +struct CustomMin +{ + __device__ __forceinline__ + float3 operator()(const float3& a, const float3& b) const { + return { min(a.x, b.x), min(a.y, b.y), min(a.z, b.z) }; + } +}; + +struct CustomMax +{ + __device__ __forceinline__ + float3 operator()(const float3& a, const float3& b) const { + return { max(a.x, b.x), max(a.y, b.y), max(a.z, b.z) }; + } +}; + +__host__ __device__ uint32_t prepMorton(uint32_t x) +{ + x = (x | (x << 16)) & 0x030000FF; + x = (x | (x << 8)) & 0x0300F00F; + x = (x | (x << 4)) & 0x030C30C3; + x = (x | (x << 2)) & 0x09249249; + return x; +} + +__host__ __device__ uint32_t coord2Morton(float3 coord, float3 minn, float3 maxx) +{ + uint32_t x = prepMorton(((coord.x - minn.x) / (maxx.x - minn.x)) * ((1 << 10) - 1)); + uint32_t y = prepMorton(((coord.y - minn.y) / (maxx.y - minn.y)) * ((1 << 10) - 1)); + uint32_t z = prepMorton(((coord.z - minn.z) / (maxx.z - minn.z)) * ((1 << 10) - 1)); + + return x | (y << 1) | (z << 2); +} + +__global__ void coord2Morton(int P, const float3* points, float3 minn, float3 maxx, uint32_t* codes) +{ + auto idx = cg::this_grid().thread_rank(); + if (idx >= P) + return; + + codes[idx] = coord2Morton(points[idx], minn, maxx); +} + +struct MinMax +{ + float3 minn; + float3 maxx; +}; + +__global__ void boxMinMax(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes) +{ + auto idx = cg::this_grid().thread_rank(); + + MinMax me; + if (idx < P) + { + me.minn = points[indices[idx]]; + me.maxx = points[indices[idx]]; + } + else + { + me.minn = { FLT_MAX, FLT_MAX, FLT_MAX }; + me.maxx = { -FLT_MAX,-FLT_MAX,-FLT_MAX }; + } + + __shared__ MinMax redResult[BOX_SIZE]; + + for (int off = BOX_SIZE / 2; off >= 1; off /= 2) + { + if (threadIdx.x < 2 * off) + redResult[threadIdx.x] = me; + __syncthreads(); + + if (threadIdx.x < off) + { + MinMax other = redResult[threadIdx.x + off]; + me.minn.x = min(me.minn.x, other.minn.x); + me.minn.y = min(me.minn.y, other.minn.y); + me.minn.z = min(me.minn.z, other.minn.z); + me.maxx.x = max(me.maxx.x, other.maxx.x); + me.maxx.y = max(me.maxx.y, other.maxx.y); + me.maxx.z = max(me.maxx.z, other.maxx.z); + } + __syncthreads(); + } + + if (threadIdx.x == 0) + boxes[blockIdx.x] = me; +} + +__device__ __host__ float distBoxPoint(const MinMax& box, const float3& p) +{ + float3 diff = { 0, 0, 0 }; + if (p.x < box.minn.x || p.x > box.maxx.x) + diff.x = min(abs(p.x - box.minn.x), abs(p.x - box.maxx.x)); + if (p.y < box.minn.y || p.y > box.maxx.y) + diff.y = min(abs(p.y - box.minn.y), abs(p.y - box.maxx.y)); + if (p.z < box.minn.z || p.z > box.maxx.z) + diff.z = min(abs(p.z - box.minn.z), abs(p.z - box.maxx.z)); + return diff.x * diff.x + diff.y * diff.y + diff.z * diff.z; +} + +template +__device__ void updateKBest(const float3& ref, const float3& point, float* knn) +{ + float3 d = { point.x - ref.x, point.y - ref.y, point.z - ref.z }; + float dist = d.x * d.x + d.y * d.y + d.z * d.z; + for (int j = 0; j < K; j++) + { + if (knn[j] > dist) + { + float t = knn[j]; + knn[j] = dist; + dist = t; + } + } +} + +__global__ void boxMeanDist(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes, float* dists) +{ + int idx = cg::this_grid().thread_rank(); + if (idx >= P) + return; + + float3 point = points[indices[idx]]; + float best[3] = { FLT_MAX, FLT_MAX, FLT_MAX }; + + for (int i = max(0, idx - 3); i <= min(P - 1, idx + 3); i++) + { + if (i == idx) + continue; + updateKBest<3>(point, points[indices[i]], best); + } + + float reject = best[2]; + best[0] = FLT_MAX; + best[1] = FLT_MAX; + best[2] = FLT_MAX; + + for (int b = 0; b < (P + BOX_SIZE - 1) / BOX_SIZE; b++) + { + MinMax box = boxes[b]; + float dist = distBoxPoint(box, point); + if (dist > reject || dist > best[2]) + continue; + + for (int i = b * BOX_SIZE; i < min(P, (b + 1) * BOX_SIZE); i++) + { + if (i == idx) + continue; + updateKBest<3>(point, points[indices[i]], best); + } + } + dists[indices[idx]] = (best[0] + best[1] + best[2]) / 3.0f; +} + +void SimpleKNN::knn(int P, float3* points, float* meanDists) +{ + float3* result; + cudaMalloc(&result, sizeof(float3)); + size_t temp_storage_bytes; + + float3 init = { 0, 0, 0 }, minn, maxx; + + cub::DeviceReduce::Reduce(nullptr, temp_storage_bytes, points, result, P, CustomMin(), init); + thrust::device_vector temp_storage(temp_storage_bytes); + + cub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMin(), init); + cudaMemcpy(&minn, result, sizeof(float3), cudaMemcpyDeviceToHost); + + cub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMax(), init); + cudaMemcpy(&maxx, result, sizeof(float3), cudaMemcpyDeviceToHost); + + thrust::device_vector morton(P); + thrust::device_vector morton_sorted(P); + coord2Morton << <(P + 255) / 256, 256 >> > (P, points, minn, maxx, morton.data().get()); + + thrust::device_vector indices(P); + thrust::sequence(indices.begin(), indices.end()); + thrust::device_vector indices_sorted(P); + + cub::DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P); + temp_storage.resize(temp_storage_bytes); + + cub::DeviceRadixSort::SortPairs(temp_storage.data().get(), temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P); + + uint32_t num_boxes = (P + BOX_SIZE - 1) / BOX_SIZE; + thrust::device_vector boxes(num_boxes); + boxMinMax << > > (P, points, indices_sorted.data().get(), boxes.data().get()); + boxMeanDist << > > (P, points, indices_sorted.data().get(), boxes.data().get(), meanDists); + + cudaFree(result); +} diff --git a/Libs/simple-knn/simple_knn.h b/Libs/simple-knn/simple_knn.h new file mode 100644 index 0000000..3fcfdb8 --- /dev/null +++ b/Libs/simple-knn/simple_knn.h @@ -0,0 +1,21 @@ +/* + * Copyright (C) 2023, Inria + * GRAPHDECO research group, https://team.inria.fr/graphdeco + * All rights reserved. + * + * This software is free for non-commercial, research and evaluation use + * under the terms of the LICENSE.md file. + * + * For inquiries contact george.drettakis@inria.fr + */ + +#ifndef SIMPLEKNN_H_INCLUDED +#define SIMPLEKNN_H_INCLUDED + +class SimpleKNN +{ +public: + static void knn(int P, float3* points, float* meanDists); +}; + +#endif \ No newline at end of file diff --git a/Libs/simple-knn/simple_knn/.gitkeep b/Libs/simple-knn/simple_knn/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/Libs/simple-knn/spatial.cu b/Libs/simple-knn/spatial.cu new file mode 100644 index 0000000..1a6a654 --- /dev/null +++ b/Libs/simple-knn/spatial.cu @@ -0,0 +1,26 @@ +/* + * Copyright (C) 2023, Inria + * GRAPHDECO research group, https://team.inria.fr/graphdeco + * All rights reserved. + * + * This software is free for non-commercial, research and evaluation use + * under the terms of the LICENSE.md file. + * + * For inquiries contact george.drettakis@inria.fr + */ + +#include "spatial.h" +#include "simple_knn.h" + +torch::Tensor +distCUDA2(const torch::Tensor& points) +{ + const int P = points.size(0); + + auto float_opts = points.options().dtype(torch::kFloat32); + torch::Tensor means = torch::full({P}, 0.0, float_opts); + + SimpleKNN::knn(P, (float3*)points.contiguous().data(), means.contiguous().data()); + + return means; +} \ No newline at end of file diff --git a/Libs/simple-knn/spatial.h b/Libs/simple-knn/spatial.h new file mode 100644 index 0000000..280c953 --- /dev/null +++ b/Libs/simple-knn/spatial.h @@ -0,0 +1,14 @@ +/* + * Copyright (C) 2023, Inria + * GRAPHDECO research group, https://team.inria.fr/graphdeco + * All rights reserved. + * + * This software is free for non-commercial, research and evaluation use + * under the terms of the LICENSE.md file. + * + * For inquiries contact george.drettakis@inria.fr + */ + +#include + +torch::Tensor distCUDA2(const torch::Tensor& points); \ No newline at end of file diff --git a/gen_3d_modules/InstantMesh/models/lrm.py b/gen_3d_modules/InstantMesh/models/lrm.py index 2128bac..88b8a77 100644 --- a/gen_3d_modules/InstantMesh/models/lrm.py +++ b/gen_3d_modules/InstantMesh/models/lrm.py @@ -15,7 +15,7 @@ import numpy as np import torch import torch.nn as nn -import mcubes +from mcubes import marching_cubes import nvdiffrast.torch as dr from einops import rearrange, repeat @@ -175,7 +175,7 @@ class InstantNeRF(nn.Module): grid_size=mesh_resolution, ) - vertices, faces = mcubes.marching_cubes( + vertices, faces = marching_cubes( grid_out['sigma'].squeeze(0).squeeze(-1).cpu().numpy(), mesh_threshold, ) diff --git a/gen_3d_modules/TripoSR/models/isosurface.py b/gen_3d_modules/TripoSR/models/isosurface.py index 39c8389..c4a9e31 100644 --- a/gen_3d_modules/TripoSR/models/isosurface.py +++ b/gen_3d_modules/TripoSR/models/isosurface.py @@ -1,9 +1,10 @@ from typing import Callable, Optional, Tuple -import numpy as np import torch import torch.nn as nn -from torchmcubes import marching_cubes +from mcubes import marching_cubes + +from mesh_processer.mesh_utils import switch_vector_axis class IsosurfaceHelper(nn.Module): @@ -42,7 +43,8 @@ class MarchingCubeHelper(IsosurfaceHelper): level: torch.FloatTensor, ) -> Tuple[torch.FloatTensor, torch.LongTensor]: level = -level.view(self.resolution, self.resolution, self.resolution) - v_pos, t_pos_idx = self.mc_func(level.detach(), 0.0) + v_pos, t_pos_idx = self.mc_func(level.detach().cpu().numpy(), 0.0) + v_pos = switch_vector_axis(v_pos, (2, 1, 0)) v_pos = v_pos[..., [2, 1, 0]] v_pos = v_pos / (self.resolution - 1.0) return v_pos, t_pos_idx diff --git a/gen_3d_modules/TripoSR/system.py b/gen_3d_modules/TripoSR/system.py index c6d42f9..8405f8e 100644 --- a/gen_3d_modules/TripoSR/system.py +++ b/gen_3d_modules/TripoSR/system.py @@ -173,7 +173,7 @@ class TSR(BaseModule): )["density_act"] v_pos, t_pos_idx = self.isosurface_helper(-(density - threshold)) v_pos = scale_tensor( - v_pos, + torch.from_numpy(v_pos).float().to(scene_codes.device), self.isosurface_helper.points_range, (-self.renderer.cfg.radius, self.renderer.cfg.radius), ) @@ -185,7 +185,7 @@ class TSR(BaseModule): )["color"] mesh = trimesh.Trimesh( vertices=v_pos.cpu().numpy(), - faces=t_pos_idx.cpu().numpy(), + faces=t_pos_idx, vertex_colors=color.cpu().numpy(), ) meshes.append(mesh) diff --git a/mesh_processer/mesh_utils.py b/mesh_processer/mesh_utils.py index 963d924..191b5cc 100644 --- a/mesh_processer/mesh_utils.py +++ b/mesh_processer/mesh_utils.py @@ -406,18 +406,18 @@ def get_target_axis_and_scale(axis_string, scale_value=1.0): return target_axis, target_scale, coordinate_invert_count -def switch_vector_axis(vector3_tensor, target_axis): +def switch_vector_axis(vector3s, target_axis): """ Example: - vector3_tensor = torch.tensor([[1, 2, 3], [3, 2, 1], [2, 3, 1]]) # shape (N, 3) + vector3s = torch.tensor([[1, 2, 3], [3, 2, 1], [2, 3, 1]]) # shape (N, 3) target_axis = (2, 0, 1) # or [2, 0, 1] - vector3_tensor[:, [0, 1, 2]] = vector3_tensor[:, target_axis] + vector3s[:, [0, 1, 2]] = vector3s[:, target_axis] # Result: tensor([[3, 1, 2], [1, 3, 2], [1, 2, 3]]) """ - vector3_tensor[:, [0, 1, 2]] = vector3_tensor[:, target_axis] - return vector3_tensor + vector3s[:, [0, 1, 2]] = vector3s[:, target_axis] + return vector3s def switch_ply_axis_and_scale(plydata, target_axis, target_scale, coordinate_invert_count): """ @@ -462,7 +462,7 @@ def switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal=Fals def marching_cubes_density_to_mesh(get_density_func, grid_size=256, S=128, density_thresh=10, decimate_target=5e4): - import mcubes + from mcubes import marching_cubes from kiui.mesh_utils import clean_mesh, decimate_mesh @@ -482,7 +482,7 @@ def marching_cubes_density_to_mesh(get_density_func, grid_size=256, S=128, densi print(f'[INFO] marching cubes thresh: {density_thresh} ({sigmas.min()} ~ {sigmas.max()})') - vertices, triangles = mcubes.marching_cubes(sigmas, density_thresh) + vertices, triangles = marching_cubes(sigmas, density_thresh) vertices = vertices / (grid_size - 1.0) * 2 - 1 # clean diff --git a/requirements_post.txt b/requirements_post.txt index 4f61218..ac5d41a 100644 --- a/requirements_post.txt +++ b/requirements_post.txt @@ -1,6 +1,6 @@ # local ./Gen_3D_Modules/TriplaneGaussian/models/snowflake/pointnet2_ops_lib -./simple-knn +./Libs/simple-knn # remote git+https://github.com/ashawkey/diff-gaussian-rasterization @@ -8,5 +8,4 @@ git+https://github.com/NVlabs/nvdiffrast/ git+https://github.com/facebookresearch/pytorch3d.git@stable git+https://github.com/rusty1s/pytorch_scatter.git git+https://github.com/ashawkey/kiuikit.git -git+https://github.com/tatsy/torchmcubes.git diff --git a/requirements_post_win_py310_cu121.txt b/requirements_post_win_py310_cu121.txt index 41166b2..ccb4b70 100644 --- a/requirements_post_win_py310_cu121.txt +++ b/requirements_post_win_py310_cu121.txt @@ -4,5 +4,4 @@ _Pre_Builds/_Wheels_win_py310_cu121/nvdiffrast-0.3.1-py3-none-any.whl _Pre_Builds/_Wheels_win_py310_cu121/pointnet2_ops-3.0.0-cp310-cp310-win_amd64.whl _Pre_Builds/_Wheels_win_py310_cu121/pytorch3d-0.7.6-cp310-cp310-win_amd64.whl _Pre_Builds/_Wheels_win_py310_cu121/simple_knn-0.0.0-cp310-cp310-win_amd64.whl -_Pre_Builds/_Wheels_win_py310_cu121/torch_scatter-2.1.2-cp310-cp310-win_amd64.whl -_Pre_Builds/_Wheels_win_py310_cu121/torchmcubes-0.1.0-cp310-cp310-win_amd64.whl \ No newline at end of file +_Pre_Builds/_Wheels_win_py310_cu121/torch_scatter-2.1.2-cp310-cp310-win_amd64.whl \ No newline at end of file diff --git a/requirements_post_win_py311_cu121.txt b/requirements_post_win_py311_cu121.txt index fac1fda..fa73cf9 100644 --- a/requirements_post_win_py311_cu121.txt +++ b/requirements_post_win_py311_cu121.txt @@ -4,5 +4,4 @@ _Pre_Builds/_Wheels_win_py311_cu121/nvdiffrast-0.3.1-py3-none-any.whl _Pre_Builds/_Wheels_win_py311_cu121/pointnet2_ops-3.0.0-cp311-cp311-win_amd64.whl _Pre_Builds/_Wheels_win_py311_cu121/pytorch3d-0.7.6-cp311-cp311-win_amd64.whl _Pre_Builds/_Wheels_win_py311_cu121/simple_knn-0.0.0-cp311-cp311-win_amd64.whl -_Pre_Builds/_Wheels_win_py311_cu121/torch_scatter-2.1.2-cp311-cp311-win_amd64.whl -_Pre_Builds/_Wheels_win_py311_cu121/torchmcubes-0.1.0-cp311-cp311-win_amd64.whl \ No newline at end of file +_Pre_Builds/_Wheels_win_py311_cu121/torch_scatter-2.1.2-cp311-cp311-win_amd64.whl \ No newline at end of file diff --git a/requirements_post_win_py312_cu121.txt b/requirements_post_win_py312_cu121.txt index b100d0c..586d15b 100644 --- a/requirements_post_win_py312_cu121.txt +++ b/requirements_post_win_py312_cu121.txt @@ -4,5 +4,4 @@ _Pre_Builds/_Wheels_win_py312_cu121/nvdiffrast-0.3.1-py3-none-any.whl _Pre_Builds/_Wheels_win_py312_cu121/pointnet2_ops-3.0.0-cp312-cp312-win_amd64.whl _Pre_Builds/_Wheels_win_py312_cu121/pytorch3d-0.7.6-cp312-cp312-win_amd64.whl _Pre_Builds/_Wheels_win_py312_cu121/simple_knn-0.0.0-cp312-cp312-win_amd64.whl -_Pre_Builds/_Wheels_win_py312_cu121/torch_scatter-2.1.2-cp312-cp312-win_amd64.whl -_Pre_Builds/_Wheels_win_py312_cu121/torchmcubes-0.1.0-cp312-cp312-win_amd64.whl \ No newline at end of file +_Pre_Builds/_Wheels_win_py312_cu121/torch_scatter-2.1.2-cp312-cp312-win_amd64.whl \ No newline at end of file