nodes/MrForExample-ComfyUI-3D-Pack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Replace torchmcubes with PyMCubes

Browse Source 
torchmcubes dependencies is archived on github, so for better maintenance, I replace it with PyMCubes inside TripoSR
This commit is contained in:
Mr. For Example
2024-07-03 17:52:40 -06:00
parent 4b22660077
commit 4e317c8d31
16 changed files with 356 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockerfile
+1 -1
View File
@@ -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
Libs/simple-knn/ext.cpp
+17
View File
@@ -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 <torch/extension.h>
#include "spatial.h"
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("distCUDA2", &distCUDA2);
}
Libs/simple-knn/setup.py
+35
View File
@@ -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
}
)
Libs/simple-knn/simple_knn.cu
+222
View File
@@ -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 <cub/cub.cuh>
#include <cub/device/device_radix_sort.cuh>
#include <vector>
#include <cuda_runtime_api.h>
#include <thrust/device_vector.h>
#include <thrust/sequence.h>
#define __CUDACC__
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include <cfloat>
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<int K>
__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<char> 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<uint32_t> morton(P);
thrust::device_vector<uint32_t> morton_sorted(P);
coord2Morton << <(P + 255) / 256, 256 >> > (P, points, minn, maxx, morton.data().get());
thrust::device_vector<uint32_t> indices(P);
thrust::sequence(indices.begin(), indices.end());
thrust::device_vector<uint32_t> 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<MinMax> boxes(num_boxes);
boxMinMax << <num_boxes, BOX_SIZE >> > (P, points, indices_sorted.data().get(), boxes.data().get());
boxMeanDist << <num_boxes, BOX_SIZE >> > (P, points, indices_sorted.data().get(), boxes.data().get(), meanDists);
cudaFree(result);
}
Libs/simple-knn/simple_knn.h
+21
View File
@@ -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
Libs/simple-knn/simple_knn/.gitkeep
View File
Libs/simple-knn/spatial.cu
+26
View File
@@ -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<float>(), means.contiguous().data<float>());
return means;
}
Libs/simple-knn/spatial.h
+14
View File
@@ -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/extension.h>
torch::Tensor distCUDA2(const torch::Tensor& points);
gen_3d_modules/InstantMesh/models/lrm.py
+2 -2
View File
@@ -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,
)
gen_3d_modules/TripoSR/models/isosurface.py
+5 -3
View File
@@ -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
gen_3d_modules/TripoSR/system.py
+2 -2
View File
@@ -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)
mesh_processer/mesh_utils.py
+7 -7
View File
@@ -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
requirements_post.txt
+1 -2
View File
@@ -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
requirements_post_win_py310_cu121.txt
+1 -2
View File
@@ -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
_Pre_Builds/_Wheels_win_py310_cu121/torch_scatter-2.1.2-cp310-cp310-win_amd64.whl
requirements_post_win_py311_cu121.txt
+1 -2
View File
@@ -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
_Pre_Builds/_Wheels_win_py311_cu121/torch_scatter-2.1.2-cp311-cp311-win_amd64.whl
requirements_post_win_py312_cu121.txt
+1 -2
View File
@@ -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
_Pre_Builds/_Wheels_win_py312_cu121/torch_scatter-2.1.2-cp312-cp312-win_amd64.whl