From df3b12e3c4e329dc5a1a974412aa1ac118586ca9 Mon Sep 17 00:00:00 2001 From: "Mr. For Example" <62230687+MrForExample@users.noreply.github.com> Date: Wed, 21 Aug 2024 16:53:15 -0600 Subject: [PATCH] Finished GPU accelerated Convert_Vertex_Color_To_Texture function --- _Pre_Builds/_Build_Scripts/auto_build_all.py | 2 +- mesh_processer/mesh.py | 19 ++- mesh_processer/mesh_utils.py | 119 ++++--------------- nodes.py | 30 ++--- 4 files changed, 48 insertions(+), 122 deletions(-) diff --git a/_Pre_Builds/_Build_Scripts/auto_build_all.py b/_Pre_Builds/_Build_Scripts/auto_build_all.py index 806e6b9..5449262 100644 --- a/_Pre_Builds/_Build_Scripts/auto_build_all.py +++ b/_Pre_Builds/_Build_Scripts/auto_build_all.py @@ -45,7 +45,7 @@ def get_dependency_dir(dependency): def setup_build_env(): # Set CMake argumens fo build packages based on CUDA - os.environ["CMAKE_ARGS"]="-DBUILD_opencv_world=ON -DWITH_CUDA=ON -DCUDA_FAST_MATH=ON -DWITH_CUBLAS=ON -DCUDA_ARCH_PTX=9.0 -DWITH_NVCUVID=ON" + os.environ["CMAKE_ARGS"] = "-DBUILD_opencv_world=ON -DWITH_CUDA=ON -DCUDA_FAST_MATH=ON -DWITH_CUBLAS=ON -DCUDA_ARCH_PTX=9.0 -DWITH_NVCUVID=ON" subprocess.run([PYTHON_PATH, "-s", "-m", "pip", "install", "-r", BUILD_REQUIREMENTS_FILE_ABS_PATH]) install_remote_packages(build_config.build_base_packages) diff --git a/mesh_processer/mesh.py b/mesh_processer/mesh.py index 1f0d259..d1247a4 100644 --- a/mesh_processer/mesh.py +++ b/mesh_processer/mesh.py @@ -540,10 +540,7 @@ class Mesh: vmapping (np.ndarray, optional): the mapping relationship from f to ft. Defaults to None. """ if vmapping is None: - ft = self.ft.view(-1).long() - f = self.f.view(-1).long() - vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device) - vmapping[ft] = f # scatter, randomly choose one if index is not unique + vmapping = self.get_default_vmapping() self.v = self.v[vmapping] if self.vc is not None: @@ -553,6 +550,20 @@ class Mesh: if self.vn is not None: self.vn = self.vn[vmapping] self.fn = self.ft + + print(f"########## !!!!!!!!!! vmapping: {vmapping.shape}; self.vt.shape: {self.vt.shape}; self.v: {self.v.shape}; self.f: {self.f.shape}") + + def get_default_vmapping(self): + """map from ft to f + + Returns: + vmapping: tensor array + """ + ft = self.ft.view(-1).long() + f = self.f.view(-1).long() + vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device) + vmapping[ft] = f # scatter, randomly choose one if index is not unique + return vmapping def to(self, device): """move all tensor attributes to device. diff --git a/mesh_processer/mesh_utils.py b/mesh_processer/mesh_utils.py index a20367e..bec4f6a 100644 --- a/mesh_processer/mesh_utils.py +++ b/mesh_processer/mesh_utils.py @@ -582,115 +582,35 @@ def K_nearest_neighbors_func( return nn[0, :, :, :], idx[0, :, :] else: return nn[0, :, :, :], idx[0, :, :], dist[0, :, :] - -def _interpolate_texture_map(mesh, texture_size: int = 256): - # Get UV coordinates and faces - if mesh.vt is None: - mesh.auto_uv() - verts_uvs = mesh.vt - verts_uvs_idx = (verts_uvs * (texture_size - 1)).to(torch.int) - faces = mesh.f - verts_colors = mesh.vc - - # Create a blank texture map - texture_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device) - - # Create a grid of UV coordinates - grid_x, grid_y = torch.meshgrid(torch.linspace(0, 1, texture_size, device=verts_colors.device), torch.linspace(0, 1, texture_size, device=verts_colors.device)) - #grid = torch.stack([grid_x, grid_y], dim=-1).reshape(-1, 2) - grid = torch.stack([grid_x, grid_y], dim=-1) - - # Function to interpolate colors within a triangle - def _interpolate_triangle(grid, uvs, colors): - # Compute barycentric coordinates - v0, v1, v2 = uvs - denom = (v1[1] - v2[1]) * (v0[0] - v2[0]) + (v2[0] - v1[0]) * (v0[1] - v2[1]) - a = ((v1[1] - v2[1]) * (grid[:, 0] - v2[0]) + (v2[0] - v1[0]) * (grid[:, 1] - v2[1])) / denom - b = ((v2[1] - v0[1]) * (grid[:, 0] - v2[0]) + (v0[0] - v2[0]) * (grid[:, 1] - v2[1])) / denom - c = 1 - a - b - - # Mask for points inside the triangle - mask = (a >= 0) & (b >= 0) & (c >= 0) - - uv_coords = (grid[mask] * (texture_size - 1)).to(torch.int) - - # Interpolate colors - interpolated_colors = a[:, None] * colors[0] + b[:, None] * colors[1] + c[:, None] * colors[2] - - #print(f"############sub_grid: {grid.shape}; a: {a.shape}; b: {b.shape}; mask: {mask.shape}; uv_coords: {grid[mask].shape}; interpolated_colors: {interpolated_colors[mask].shape}") - - #return grid[mask], interpolated_colors[mask] - return uv_coords, interpolated_colors[mask] - - # Function to interpolate colors within a triangle - def interpolate_triangle(grid, uvs_idx, uvs, colors): - u_idx, v_idx = uvs_idx[:, 0], uvs_idx[:, 1] - sub_grid = grid[torch.min(u_idx):torch.max(u_idx)+1, torch.min(v_idx):torch.max(v_idx)+1] - #sub_grid = grid - # Compute barycentric coordinates - v0, v1, v2 = uvs - denom = (v1[1] - v2[1]) * (v0[0] - v2[0]) + (v2[0] - v1[0]) * (v0[1] - v2[1]) - a = ((v1[1] - v2[1]) * (sub_grid[:, :, 0] - v2[0]) + (v2[0] - v1[0]) * (sub_grid[:, :, 1] - v2[1])) / denom - b = ((v2[1] - v0[1]) * (sub_grid[:, :, 0] - v2[0]) + (v0[0] - v2[0]) * (sub_grid[:, :, 1] - v2[1])) / denom - c = 1 - a - b - - # Mask for points inside the triangle - mask = (a >= 0) & (b >= 0) & (c >= 0) - - # Get texture coordinates - #uv_coords = (sub_grid[mask][[1,0]] * (texture_size - 1)).to(torch.int) - uv_coords = (sub_grid[mask] * (texture_size - 1)).to(torch.int) - - # Interpolate colors - interpolated_colors = (a[:, :, None] * colors[0] + b[:, :, None] * colors[1] + c[:, :, None] * colors[2])[mask] - - #print(f"############sub_grid: {sub_grid.shape}; a: {a.shape}; b: {b.shape}; mask: {mask.shape}; uv_coords: {uv_coords.shape}; interpolated_colors: {interpolated_colors.shape}") - - return uv_coords, interpolated_colors - - # Map vertex colors to the UV coordinates - import comfy.utils - comfy_pbar = comfy.utils.ProgressBar(faces.shape[0]) - for i, face in enumerate(faces): - uvs = verts_uvs[face] - uvs_idx = verts_uvs_idx[face] - colors = verts_colors[face] - uv_coords, interpolated_colors = interpolate_triangle(grid, uvs_idx, uvs, colors) - #uv_coords, interpolated_colors = interpolate_triangle(grid, uvs, colors) - - # Map interpolated colors to the texture map - #for uv, color in zip(uv_coords, interpolated_colors): - # x, y = int(uv[0] * (texture_size - 1)), int(uv[1] * (texture_size - 1)) - # texture_map[y, x] = color - - texture_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_colors - - comfy_pbar.update_absolute(i + 1) - - return texture_map def interpolate_texture_map(mesh, texture_size: int = 256, batch_size: int = 64): # Get UV coordinates and faces if mesh.vt is None: mesh.auto_uv() - verts_uvs = mesh.vt - verts_uvs_idx = (mesh.vt * (texture_size - 1)).to(torch.int) - faces = mesh.f - verts_colors = mesh.vc + + # Get Faces on UV + texture_size_minus_one = texture_size - 1 + faces = mesh.ft + verts_uvs = mesh.vt * texture_size_minus_one + verts_uvs_idx = verts_uvs.to(torch.long) + + # Get vertex colors + vmapping = mesh.get_default_vmapping() + verts_colors = mesh.vc[vmapping] # Create a blank texture map texture_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device) # Create a grid of UV coordinates - grid_x, grid_y = torch.meshgrid(torch.linspace(0, 1, texture_size, device=verts_colors.device), torch.linspace(0, 1, texture_size, device=verts_colors.device)) + grid_x, grid_y = torch.meshgrid(torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device), torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device)) grid = torch.stack([grid_x, grid_y], dim=-1) - batch_idx_range = torch.arange(0, batch_size**2, 1, device=verts_colors.device, dtype=torch.int) + batch_idx_range = torch.arange(0, batch_size**2, 1, device=verts_colors.device, dtype=torch.long) - # Divide faces into squre batchs all_face_verts = verts_uvs_idx[faces] - for uv_i in range(0, texture_size-1, batch_size): - for uv_j in range(0, texture_size-1, batch_size): + for uv_i in range(0, texture_size_minus_one, batch_size): + for uv_j in range(0, texture_size_minus_one, batch_size): + # Divide faces into squre batchs faces_mask = (uv_i <= all_face_verts[:, :, 0]) & (all_face_verts[:, :, 0] <= uv_i + batch_size) & (uv_j <= all_face_verts[:, :, 1]) & (all_face_verts[:, :, 1] <= uv_j + batch_size) faces_mask, _ = torch.max(faces_mask, dim=1) faces_mask = faces_mask.nonzero().squeeze(-1) @@ -698,14 +618,15 @@ def interpolate_texture_map(mesh, texture_size: int = 256, batch_size: int = 64) uvs = verts_uvs[batch_faces] colors = verts_colors[batch_faces] - + v0, v1, v2 = uvs[:, 0], uvs[:, 1], uvs[:, 2] v0_0, v0_1 = v0[:, 0], v0[:, 1] v1_0, v1_1 = v1[:, 0], v1[:, 1] v2_0, v2_1 = v2[:, 0], v2[:, 1] sub_grid = grid[uv_i:uv_i + batch_size, uv_j:uv_j + batch_size].reshape(-1, 2).unsqueeze(1) - + + # Compute barycentric coordinates for the batch of faces denom = (v1_1 - v2_1) * (v0_0 - v2_0) + (v2_0 - v1_0) * (v0_1 - v2_1) a = ((v1_1 - v2_1) * (sub_grid[:, :, 0] - v2_0) + (v2_0 - v1_0) * (sub_grid[:, :, 1] - v2_1)) / denom b = ((v2_1 - v0_1) * (sub_grid[:, :, 0] - v2_0) + (v0_0 - v2_0) * (sub_grid[:, :, 1] - v2_1)) / denom @@ -714,10 +635,10 @@ def interpolate_texture_map(mesh, texture_size: int = 256, batch_size: int = 64) # Mask for points inside the triangle mask = (a >= 0) & (b >= 0) & (c >= 0) mask, mask_idx = torch.max(mask, dim=1) - + # Get texture coordinates sub_grid = sub_grid.squeeze(1) - uv_coords = (sub_grid[mask] * (texture_size - 1)).to(torch.int) + uv_coords = sub_grid[mask].to(torch.long) # Interpolate colors a, b, c = a[batch_idx_range, mask_idx][mask], b[batch_idx_range, mask_idx][mask], c[batch_idx_range, mask_idx][mask] diff --git a/nodes.py b/nodes.py index 426f0b1..d0e5196 100644 --- a/nodes.py +++ b/nodes.py @@ -43,7 +43,6 @@ from .mesh_processer.mesh_utils import ( calculate_max_sh_degree_from_gs_ply, marching_cubes_density_to_mesh, color_func_to_albedo, - K_nearest_neighbors_func, interpolate_texture_map, ) @@ -3018,10 +3017,15 @@ class ExplicitTarget_Color_Projection: new_meshes = multiview_color_projection(meshes, pil_image_list, weights=weights, resolution=projection_resolution, device=DEVICE, complete_unseen=complete_unseen_rgb, confidence_threshold=confidence_threshold, cameras_list=cam_list) vertices, faces, vertex_colors = from_py3d_mesh(new_meshes) - vertices = vertices / 2 * 1.35 - mesh = Mesh(v=vertices, f=faces, vc=vertex_colors, device=DEVICE) - mesh.auto_normal() + print(f"###############faces: {faces.shape} \nmesh.f: {mesh.f.shape} \nvertices: {vertices.shape} \nmesh.v: {mesh.v.shape} \nmesh.vt: {mesh.vt.shape}") + + mesh = Mesh(v=vertices, f=faces, + vn=None if mesh.vn is None else mesh.vn.clone(), fn=None if mesh.fn is None else mesh.fn.clone(), + vt=None if mesh.vt is None else mesh.vt.clone(), ft=None if mesh.ft is None else mesh.ft.clone(), + vc=vertex_colors, device=DEVICE) + if mesh.vn is None: + mesh.auto_normal() return (mesh,) class Convert_Vertex_Color_To_Texture: @@ -3031,8 +3035,7 @@ class Convert_Vertex_Color_To_Texture: "required": { "mesh": ("MESH",), "texture_resolution": ("INT", {"default": 1024, "min": 128, "max": 8192}), - "K_nearest_neighbors": ("INT", {"default": 3, "min": 1, "max": 0xffffffffffffffff}), - "force_cuda_rast": ("BOOLEAN", {"default": False}), + "batch_size": ("INT", {"default": 128, "min": 1, "max": 0xffffffffffffffff}), }, } @@ -3045,25 +3048,16 @@ class Convert_Vertex_Color_To_Texture: FUNCTION = "run_convert_func" CATEGORY = "Comfy3D/Algorithm" - def run_convert_func(self, mesh, texture_resolution, K_nearest_neighbors, force_cuda_rast): + def run_convert_func(self, mesh, texture_resolution, batch_size): if mesh.vc is not None: - #self.v = mesh.v.detach() - #self.vc = mesh.vc.detach() - #self.K = K_nearest_neighbors - #mesh.albedo = color_func_to_albedo(mesh, self.get_color_func, texture_resolution, device=DEVICE, force_cuda_rast=force_cuda_rast) - - mesh.albedo = interpolate_texture_map(mesh, texture_resolution) + albedo_img = interpolate_texture_map(mesh, texture_resolution, batch_size) + mesh.albedo = troch_image_dilate(albedo_img) else: cstr(f"[{self.__class__.__name__}] skip this node since there is no vertex color found in mesh").msg.print() return (mesh,) - def get_color_func(self, xs): - _, idx, dist = K_nearest_neighbors_func(self.v, self.K, xs, return_dist=True, skip_first_index=False) - colors = torch.sum(self.vc[idx] * (dist / torch.sum(dist, dim=-1, keepdim=True)).unsqueeze(-1), dim=1) - return colors - class Load_CharacterGen_MVDiffusion_Model: checkpoints_dir = "CharacterGen" default_repo_id = "zjpshadow/CharacterGen"