From 361fa98202e8521aafdf8abc0dda7e2c35cf1463 Mon Sep 17 00:00:00 2001 From: kijai Date: Fri, 3 Jul 2026 11:42:59 +0300 Subject: [PATCH] Add RenderMesh, cleanup --- .../mesh3d/postprocess/qem_decimate.py | 17 +- comfy_extras/mesh3d/uv_unwrap/pack.py | 6 +- comfy_extras/mesh3d/uv_unwrap/segment.py | 2 +- comfy_extras/nodes_mesh_postprocess.py | 709 +++++++++--------- comfy_extras/nodes_save_3d.py | 224 ++---- 5 files changed, 440 insertions(+), 518 deletions(-) diff --git a/comfy_extras/mesh3d/postprocess/qem_decimate.py b/comfy_extras/mesh3d/postprocess/qem_decimate.py index 9ad28e769..157364aaa 100644 --- a/comfy_extras/mesh3d/postprocess/qem_decimate.py +++ b/comfy_extras/mesh3d/postprocess/qem_decimate.py @@ -13,7 +13,6 @@ from dataclasses import dataclass from typing import Optional, Tuple import math -import time as _time import numpy as _np import torch @@ -860,14 +859,13 @@ class CleanStats: duplicate_faces: int = 0 # same vertex-set removed unused_verts: int = 0 # verts not in any face removed components_dropped: int = 0 # disconnected components below threshold - seconds: float = 0.0 def __str__(self): return (f"clean: in={self.in_verts}v/{self.in_faces}f -> " f"out={self.out_verts}v/{self.out_faces}f " f"(welded {self.welded_verts}v, degen {self.degenerate_faces}f, " f"dup {self.duplicate_faces}f, unused {self.unused_verts}v, " - f"comps {self.components_dropped}) {self.seconds*1000:.1f}ms") + f"comps {self.components_dropped})") def _weld_vertices( @@ -1150,7 +1148,6 @@ def clean_mesh( ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], CleanStats]: """Mesh hygiene pipeline; preserves per-vertex attributes through welding. Returns (v, f, colors, normals, stats).""" stats = CleanStats(in_verts=verts.shape[0], in_faces=faces.shape[0]) - t0 = _time.perf_counter() v = verts f = faces.long() if faces.numel() > 0 else faces c = colors @@ -1192,7 +1189,6 @@ def clean_mesh( stats.out_verts = v.shape[0] stats.out_faces = f.shape[0] - stats.seconds = _time.perf_counter() - t0 # materialize tensor-scalar counts to plain ints once at exit for field in ("welded_verts", "degenerate_faces", "duplicate_faces", "unused_verts", "components_dropped"): @@ -1210,8 +1206,6 @@ class SimplifyStats: output_faces: int = 0 iterations: int = 0 total_collapses: int = 0 - seconds: float = 0.0 - peak_mem_mb: float = 0.0 def qem_simplify( @@ -1258,11 +1252,6 @@ def qem_simplify( (normals_w.to(in_n_dtype) if normals_w is not None else None), \ stats - if device.type == "cuda": - torch.cuda.synchronize(device) - torch.cuda.reset_peak_memory_stats(device) - t0 = _time.perf_counter() - v_alive = torch.ones(num_verts, dtype=torch.bool, device=device) f_alive = torch.ones(num_faces, dtype=torch.bool, device=device) @@ -1600,10 +1589,6 @@ def qem_simplify( wrong = (fn * ref).sum(dim=-1) < 0 final_f[wrong] = final_f[wrong][:, [0, 2, 1]] - if device.type == "cuda": - torch.cuda.synchronize(device) - stats.peak_mem_mb = torch.cuda.max_memory_allocated(device) / (1024 * 1024) - stats.seconds = _time.perf_counter() - t0 stats.iterations = iteration stats.total_collapses = total_collapses stats.output_verts = final_v.shape[0] diff --git a/comfy_extras/mesh3d/uv_unwrap/pack.py b/comfy_extras/mesh3d/uv_unwrap/pack.py index 64b537c71..5b49b613f 100644 --- a/comfy_extras/mesh3d/uv_unwrap/pack.py +++ b/comfy_extras/mesh3d/uv_unwrap/pack.py @@ -9,6 +9,8 @@ import numpy as np import torch from torch import Tensor +import comfy.model_management + try: from numba import njit as _njit _HAVE_NUMBA_PACK = True @@ -361,7 +363,7 @@ def _chart_perimeter(uvs: np.ndarray, faces: np.ndarray) -> float: return float(_chart_perimeter_jit(uvs.astype(np.float64), faces.astype(np.int64), V)) -# ---- Torch fallback (used when numba is unavailable; runs on GPU if present) ---- +# Torch fallback (used when numba is unavailable; runs on GPU if present) def _dilate_local(x: Tensor, p: int) -> Tensor: """4-connectivity dilation by p, applied per-image over a batch of (cnt,g,g) bitmaps. @@ -512,7 +514,7 @@ def _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces, n = len(chart_uvs) if n == 0: return [], 1, 1 - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = comfy.model_management.get_torch_device() ang = torch.linspace(0.0, math.pi / 2.0, 37, device=device)[:-1] cos_a, sin_a = ang.cos(), ang.sin() diff --git a/comfy_extras/mesh3d/uv_unwrap/segment.py b/comfy_extras/mesh3d/uv_unwrap/segment.py index 73c4f2810..1983ccca8 100644 --- a/comfy_extras/mesh3d/uv_unwrap/segment.py +++ b/comfy_extras/mesh3d/uv_unwrap/segment.py @@ -324,7 +324,7 @@ def segment_charts( return _renumber(face_chart, device) -# ---- Parallel edge-collapse (PEC) chart clustering (CUDA) ---- +# Parallel edge-collapse (PEC) chart clustering (GPU) def _combine_normal_cones( axis_a: Tensor, half_a: Tensor, axis_b: Tensor, half_b: Tensor, diff --git a/comfy_extras/nodes_mesh_postprocess.py b/comfy_extras/nodes_mesh_postprocess.py index 924f1a0e6..b084fd908 100644 --- a/comfy_extras/nodes_mesh_postprocess.py +++ b/comfy_extras/nodes_mesh_postprocess.py @@ -2,13 +2,14 @@ import torch import numpy as np import math from typing_extensions import override -from comfy_api.latest import ComfyExtension, IO, Types, io +from comfy_api.latest import ComfyExtension, IO, Types import copy import comfy.utils import comfy.model_management from server import PromptServer from comfy_extras.mesh3d.postprocess.qem_decimate import QEMConfig, qem_decimate_simplify, qem_cluster_decimate, _compute_vertex_normals from comfy_extras.mesh3d.postprocess.remesh import remesh_narrow_band_dc, _point_tri_closest +from comfy_extras.nodes_save_3d import get_mesh_batch_item, pack_variable_mesh_batch from comfy_extras.mesh3d.uv_unwrap import mesh as _uv_mesh from comfy_extras.mesh3d.uv_unwrap import segment as _uv_seg from comfy_extras.mesh3d.uv_unwrap import parameterize as _uv_param @@ -20,57 +21,6 @@ from scipy.sparse.csgraph import connected_components from scipy.spatial import cKDTree import scipy.ndimage as ndi -MeshCameras = io.Custom("MESH_CAMERAS") # carries the camera set from RenderMeshViews → BakeViewsToTexture - -def get_mesh_batch_item(mesh, index): - if hasattr(mesh, "vertex_counts") and mesh.vertex_counts is not None: - vertex_count = int(mesh.vertex_counts[index].item()) - face_count = int(mesh.face_counts[index].item()) - vertices = mesh.vertices[index, :vertex_count] - faces = mesh.faces[index, :face_count] - colors = None - if hasattr(mesh, "colors") and mesh.colors is not None: - if hasattr(mesh, "color_counts") and mesh.color_counts is not None: - color_count = int(mesh.color_counts[index].item()) - colors = mesh.colors[index, :color_count] - else: - colors = mesh.colors[index, :vertex_count] - return vertices, faces, colors - - colors = None - if hasattr(mesh, "colors") and mesh.colors is not None: - colors = mesh.colors[index] - return mesh.vertices[index], mesh.faces[index], colors - -def pack_variable_mesh_batch(vertices, faces, colors=None): - batch_size = len(vertices) - max_vertices = max(v.shape[0] for v in vertices) - max_faces = max(f.shape[0] for f in faces) - - packed_vertices = vertices[0].new_zeros((batch_size, max_vertices, vertices[0].shape[1])) - packed_faces = faces[0].new_zeros((batch_size, max_faces, faces[0].shape[1])) - vertex_counts = torch.tensor([v.shape[0] for v in vertices], device=vertices[0].device, dtype=torch.int64) - face_counts = torch.tensor([f.shape[0] for f in faces], device=faces[0].device, dtype=torch.int64) - - for i, (v, f) in enumerate(zip(vertices, faces)): - packed_vertices[i, :v.shape[0]] = v - packed_faces[i, :f.shape[0]] = f - - mesh = Types.MESH(packed_vertices, packed_faces) - mesh.vertex_counts = vertex_counts - mesh.face_counts = face_counts - - if colors is not None: - max_colors = max(c.shape[0] for c in colors) - packed_colors = colors[0].new_zeros((batch_size, max_colors, colors[0].shape[1])) - color_counts = torch.tensor([c.shape[0] for c in colors], device=colors[0].device, dtype=torch.int64) - for i, c in enumerate(colors): - packed_colors[i, :c.shape[0]] = c - mesh.vertex_colors = packed_colors - mesh.color_counts = color_counts - - return mesh - def paint_mesh_with_voxels(mesh, voxel_coords, voxel_colors, resolution): """Paint a mesh using nearest-neighbor colors from a sparse voxel field.""" @@ -107,13 +57,21 @@ def paint_mesh_with_voxels(mesh, voxel_coords, voxel_colors, resolution): return out_mesh + +def paint_mesh_default_colors(mesh): + out_mesh = copy.copy(mesh) + vertex_count = mesh.vertices.shape[1] + out_mesh.vertex_colors = mesh.vertices.new_zeros((1, vertex_count, 3)) + return out_mesh + + class PaintMesh(IO.ComfyNode): @classmethod def define_schema(cls): return IO.Schema( node_id="PaintMesh", display_name="Paint Mesh", - category="latent/3d", + category="3d/mesh", description="Paints each mesh vertex with its nearest voxel color from the input voxel field.", inputs=[ IO.Mesh.Input("mesh"), @@ -146,7 +104,7 @@ class PaintMesh(IO.ComfyNode): sel = batch_idx == i item_coords = voxel_coords[sel] item_colors = colors[sel] - item_vertices, item_faces, _ = get_mesh_batch_item(mesh, i) + item_vertices, item_faces, *_ = get_mesh_batch_item(mesh, i) item_mesh = Types.MESH(vertices=item_vertices.unsqueeze(0), faces=item_faces.unsqueeze(0)) if item_coords.shape[0] == 0: @@ -445,14 +403,15 @@ def _sample_voxel_attrs_per_texel(position_map, mask, voxel_coords, voxel_colors origin = np.array([-0.5, -0.5, -0.5], dtype=np.float32) voxel_size = 1.0 / float(resolution) coords_np = voxel_coords.detach().cpu().numpy() - # Cell-CENTER convention (+0.5 voxel) — same world mapping as the GPU paths; this - # cKDTree only serves the rare non-CUDA nearest fallback. + # Cell-CENTER convention (+0.5 voxel) — same world mapping as the sampling paths; these + # voxel centres serve the rare cKDTree nearest fallback below. voxel_pos = (coords_np.astype(np.float32) + 0.5) * voxel_size + origin valid_positions = position_map[mask] def _nearest(query): - # GPU grid scan + small-N brute tail. Large straggler counts (coarse mesh) and - # non-CUDA come back unfound → resolve with one cKDTree (build amortizes over N). + # Grid scan + small-N brute tail (on the compute device). Only a large count of far + # stragglers (coarse mesh, or a surface off the voxel shell) is left unfound → resolve + # those with one cKDTree, since GPU brute force is O(N·M) and blows up at large N. vals, found = _nearest_voxel_sample_gpu(query, coords_np, color_np, resolution) if not found.all(): tree = cKDTree(voxel_pos) @@ -860,12 +819,9 @@ def _bake_ambient_occlusion(high_v, high_f, low_v_np, low_f_np, low_uv_np, low_n S = int(num_samples) if ray_chunk is None: # ~376 B/ray (int32 stack max_stack*4 + a few [N,3] ray buffers); spend a quarter of free - # VRAM. Speed saturates around 4M rays/chunk, so cap there (≈2 GB peak) rather than grow - # memory for no gain; floor keeps tiny GPUs from thrashing into too many chunks. - try: - free = torch.cuda.mem_get_info(dev)[0] if dev.type == "cuda" else (2 << 30) - except RuntimeError: - free = 2 << 30 + # device memory. Speed saturates around 4M rays/chunk, so cap there (≈2 GB peak) rather than + # grow memory for no gain; floor keeps tiny GPUs from thrashing into too many chunks. + free = comfy.model_management.get_free_memory(dev) ray_chunk = int(min(1 << 22, max(1 << 20, (free * 0.25) / (num_samples * 4 + 200)))) face_idx, bary_uv, mask = _rasterize_uv_barycentric(low_f_np, low_uv_np, resolution) if not mask.any(): @@ -924,6 +880,178 @@ def _bake_ambient_occlusion(high_v, high_f, low_v_np, low_f_np, low_uv_np, low_n return _jfa_fill_gpu(np.ascontiguousarray(out3, dtype=np.float32), mask.cpu().numpy()) +def _camera_basis(eye, center, up_hint): + """Forward/right/up for a camera at `eye` looking at `center` (each [3]).""" + f = torch.nn.functional.normalize(center - eye, dim=-1, eps=1e-6) + # Fall back to +Z up when looking near-vertical (f ∥ +Y would give a degenerate right vector). + up = up_hint if float(torch.abs((f * up_hint).sum())) < 0.99 else torch.tensor([0.0, 0.0, 1.0], device=f.device) + r = torch.nn.functional.normalize(torch.cross(f, up, dim=-1), dim=-1, eps=1e-6) + return f, r, torch.cross(r, f, dim=-1) + + +def _sample_image01(img_hwc, uv01): + """Bilinear-sample img [H,W,C] at uv01 [K,2] in [0,1] (u=x/col, v=y/row). Returns [K,C].""" + g = (uv01 * 2.0 - 1.0).view(1, 1, -1, 2) + s = torch.nn.functional.grid_sample(img_hwc.permute(2, 0, 1)[None].float(), g, mode="bilinear", align_corners=False, padding_mode="border") + return s[0, :, 0, :].t() + + +def _render_view(tri, bvh, uv, faces, texture_hwc, eye, f, r, u, fov, H, W, ray_chunk=1 << 22, + vertex_colors=None, vertex_normals=None, render_normal=False): + """Ray-cast render: per pixel, nearest-hit triangle → colour it. With `render_normal`, output the + view-space normal (OpenGL: x=right, y=up, z=toward camera; smooth `vertex_normals` if given, else + the face normal). Otherwise colour source in order: `texture_hwc` (sampled via interpolated UVs), + else `vertex_colors` (barycentric), else neutral clay shaded by facing angle. Returns + (img, hit_mask, depth).""" + dev = tri.device + ys = 1.0 - (torch.arange(H, device=dev, dtype=torch.float32) + 0.5) / H * 2.0 # row 0 = +up + xs = (torch.arange(W, device=dev, dtype=torch.float32) + 0.5) / W * 2.0 - 1.0 + gy, gx = torch.meshgrid(ys, xs, indexing="ij") + tn = math.tan(0.5 * fov) + aspect = W / H + d = torch.nn.functional.normalize( + (r * (gx * tn * aspect)[..., None] + u * (gy * tn)[..., None] + f).reshape(-1, 3), dim=-1, eps=1e-6) + o = eye[None, :].expand(H * W, 3) + img = torch.zeros((H * W, 3), device=dev) + depth = torch.full((H * W,), float("inf"), device=dev) + hit_all = torch.zeros(H * W, dtype=torch.bool, device=dev) + for s in range(0, H * W, ray_chunk): + e = min(s + ray_chunk, H * W) + t_hit, face, hit = _closest_hit_rays_bvh(o[s:e], d[s:e], tri, bvh, tmin=1e-5, tmax=1e30) + if bool(hit.any()): + fh = face[hit].clamp_min(0) + P = o[s:e][hit] + t_hit[hit, None] * d[s:e][hit] + bary = _barycentric(P, tri[fh]) + local = torch.zeros((e - s, 3), device=dev) + if render_normal: + if vertex_normals is not None: + nrm = torch.nn.functional.normalize( + (bary[:, :, None] * vertex_normals[faces[fh]]).sum(1), dim=-1, eps=1e-6) + else: # face normal, oriented toward camera + nrm = torch.nn.functional.normalize( + torch.cross(tri[fh][:, 1] - tri[fh][:, 0], tri[fh][:, 2] - tri[fh][:, 0], dim=-1), + dim=-1, eps=1e-6) + nrm = torch.where((nrm * -d[s:e][hit]).sum(-1, keepdim=True) < 0, -nrm, nrm) + nv = torch.stack([(nrm * r).sum(-1), (nrm * u).sum(-1), (nrm * -f).sum(-1)], dim=-1) + local[hit] = (nv * 0.5 + 0.5).clamp(0.0, 1.0) # view-space OpenGL normal encode + elif texture_hwc is not None and uv is not None: + uvh = (bary[:, :, None] * uv[faces[fh]]).sum(1) + local[hit] = _sample_image01(texture_hwc, uvh) + elif vertex_colors is not None: + local[hit] = (bary[:, :, None] * vertex_colors[faces[fh]]).sum(1) + else: + # Neutral clay, headlight-shaded (|n·view|) so silhouette-plus-form reads, not a flat blob. + fn = torch.nn.functional.normalize( + torch.cross(tri[fh][:, 1] - tri[fh][:, 0], tri[fh][:, 2] - tri[fh][:, 0], dim=-1), + dim=-1, eps=1e-6) + ndl = (fn * -d[s:e][hit]).sum(-1).abs().clamp(0.15, 1.0) + local[hit] = torch.tensor([0.72, 0.72, 0.72], device=dev) * ndl[:, None] + img[s:e] = local + dloc = torch.full((e - s,), float("inf"), device=dev) + dloc[hit] = t_hit[hit] + depth[s:e] = dloc + hit_all[s:e] = hit + img = img.reshape(H, W, 3) + depth = depth.reshape(H, W) + hit_all = hit_all.reshape(H, W) + # Dilate the object color into the background so bilinear sampling near the silhouette doesn't + # bleed black (a cross-view seam source) — and gives the upscaler a coherent, edge-free image. + if bool(hit_all.any()) and not bool(hit_all.all()): + img = torch.from_numpy(_jfa_fill_gpu(img.cpu().numpy(), hit_all.cpu().numpy())).to(dev) + return img, hit_all, depth + + +def _smooth_vertex_normals(vertices_np, faces_np, weld=True): + """Area-weighted per-vertex normals (unit length), fully smooth, no vertex splitting.""" + tris = vertices_np[faces_np] # (M, 3, 3) + face_n = np.cross(tris[:, 1] - tris[:, 0], tris[:, 2] - tris[:, 0]) + if weld and vertices_np.shape[0]: + # Group coincident positions (quantized to ~1e-5 of the bbox) into one shared normal. + lo = vertices_np.min(0) + inv_tol = 1.0 / (max(float((vertices_np.max(0) - lo).max()), 1e-9) * 1e-5) + q = np.round((vertices_np - lo) * inv_tol).astype(np.int64) + _, group = np.unique(q, axis=0, return_inverse=True) + acc = np.zeros((int(group.max()) + 1, 3), dtype=np.float64) + for k in range(3): + np.add.at(acc, group[faces_np[:, k]], face_n) + normals = acc[group] # welded normal back to each vertex + else: + normals = np.zeros((vertices_np.shape[0], 3), dtype=np.float64) + for k in range(3): + np.add.at(normals, faces_np[:, k], face_n) + lens = np.linalg.norm(normals, axis=1, keepdims=True) + normals /= np.where(lens > 1e-12, lens, 1.0) + return normals.astype(np.float32) + + +def _compute_vertex_face_normals(vertices_np, faces_np, crease_angle=None): + """Compute per-vertex normals, returning (vertices, faces_uint32, normals, remap). + + crease_angle is None (or >= 180) -> fully smooth normals; vertices/faces are + returned unchanged and remap is None. + + Otherwise vertices are split along edges whose dihedral angle exceeds + crease_angle (degrees) so hard creases stay sharp while smooth regions still + interpolate. remap maps each output vertex back to its source index, so the + caller can duplicate any per-vertex attributes (uvs / colors) to match.""" + faces_i = faces_np.astype(np.int64) + if crease_angle is None or crease_angle >= 180.0: + return (vertices_np, faces_i.astype(np.uint32), + _smooth_vertex_normals(vertices_np, faces_i), None) + + M = faces_i.shape[0] + tris = vertices_np[faces_i] + face_n = np.cross(tris[:, 1] - tris[:, 0], tris[:, 2] - tris[:, 0]) + areas = np.linalg.norm(face_n, axis=1, keepdims=True) + face_unit = face_n / np.where(areas > 1e-12, areas, 1.0) + cos_thresh = math.cos(math.radians(crease_angle)) + + # Union faces that share an edge whose dihedral angle is below the crease + # threshold; each connected component becomes one smoothing group. + parent = list(range(M)) + + def find(x): + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + edge_faces = {} + for fi in range(M): + a, b, c = int(faces_i[fi, 0]), int(faces_i[fi, 1]), int(faces_i[fi, 2]) + for u, v in ((a, b), (b, c), (c, a)): + edge_faces.setdefault((u, v) if u < v else (v, u), []).append(fi) + for fl in edge_faces.values(): + if len(fl) == 2 and float(np.dot(face_unit[fl[0]], face_unit[fl[1]])) >= cos_thresh: + ra, rb = find(fl[0]), find(fl[1]) + if ra != rb: + parent[ra] = rb + + # Emit one output vertex per (original vertex, smoothing group) pair. + new_index = {} + remap = [] + out_faces = np.empty((M, 3), dtype=np.int64) + for fi in range(M): + g = find(fi) + for k in range(3): + ov = int(faces_i[fi, k]) + key = (ov, g) + ni = new_index.get(key) + if ni is None: + ni = len(remap) + new_index[key] = ni + remap.append(ov) + out_faces[fi, k] = ni + + remap = np.asarray(remap, dtype=np.int64) + normals = np.zeros((remap.shape[0], 3), dtype=np.float64) + for k in range(3): + np.add.at(normals, out_faces[:, k], face_n) + lens = np.linalg.norm(normals, axis=1, keepdims=True) + normals /= np.where(lens > 1e-12, lens, 1.0) + return (vertices_np[remap], out_faces.astype(np.uint32), normals.astype(np.float32), remap) + + def _compute_vertex_tangents(verts, faces, uvs, normals): """Per-vertex tangents (Lengyel) orthonormalized against `normals`. Returns [N,4]: unit tangent xyz + handedness w (the bitangent is w * cross(N, T)). Pure torch.""" @@ -1200,7 +1328,7 @@ class BakeTextureFromVoxel(IO.ComfyNode): return IO.Schema( node_id="BakeTextureFromVoxel", display_name="Bake Texture From Voxel", - category="latent/3d", + category="3d/mesh/texturing", description=( "Bakes PBR textures onto the mesh's existing UV layout (trilinear-sample the " "sparse voxel volume). Does NOT unwrap — connect a UV unwrap node upstream. Outputs " @@ -1249,7 +1377,7 @@ class BakeTextureFromVoxel(IO.ComfyNode): sel = batch_idx == i item_coords = voxel_xyz[sel] item_colors = colors[sel] - v_i, f_i, _ = get_mesh_batch_item(mesh, i) + v_i, f_i, *_ = get_mesh_batch_item(mesh, i) if item_coords.shape[0] == 0 or f_i.numel() == 0: logging.warning(f"BakeTextureFromVoxel: skipping batch {i} (empty voxel/mesh)") pbar.update(5) @@ -1257,7 +1385,7 @@ class BakeTextureFromVoxel(IO.ComfyNode): ev_i = mesh_uvs[i, :v_i.shape[0]] ref_i = None if reference_mesh is not None: - rv_i, rf_i, _ = get_mesh_batch_item(reference_mesh, i) + rv_i, rf_i, *_ = get_mesh_batch_item(reference_mesh, i) ref_i = (rv_i, rf_i) _bv, _bf, _bu, bt, bmr = bake_texture_from_voxel_fn( v_i, f_i, item_coords, item_colors, @@ -1303,7 +1431,7 @@ class MeshTextureToImage(IO.ComfyNode): return IO.Schema( node_id="MeshTextureToImage", display_name="Mesh Texture to Image", - category="latent/3d", + category="3d/mesh/texturing", description=( "Extracts a mesh's baked textures as individual IMAGEs: base_color, metallic, " "roughness, occlusion and normal_map. Channels with nothing baked come back " @@ -1361,7 +1489,7 @@ class ApplyTextureToMesh(IO.ComfyNode): return IO.Schema( node_id="ApplyTextureToMesh", display_name="Apply Texture to Mesh", - category="latent/3d", + category="3d/mesh/texturing", description=( "Attaches baked texture IMAGEs to a mesh's UV layout for SaveGLB. Feed the SAME mesh you baked" ), @@ -1388,7 +1516,7 @@ class ApplyTextureToMesh(IO.ComfyNode): if mesh_uvs.ndim == 3: new_uvs = mesh_uvs.clone() for i in range(mesh_uvs.shape[0]): - v_i, _f_i, _ = get_mesh_batch_item(mesh, i) + v_i, _f_i, *_ = get_mesh_batch_item(mesh, i) n = v_i.shape[0] norm = _normalize_uvs_to_unit(mesh_uvs[i, :n].detach().cpu().numpy()) new_uvs[i, :n] = torch.from_numpy(norm).to(new_uvs) @@ -1410,7 +1538,7 @@ class ApplyTextureToMesh(IO.ComfyNode): tangents_padded = torch.zeros((B, Nmax, 4), dtype=torch.float32) normals_padded = torch.zeros((B, Nmax, 3), dtype=torch.float32) for i in range(B): - v_i, f_i, _ = get_mesh_batch_item(mesh, i) + v_i, f_i, *_ = get_mesh_batch_item(mesh, i) n = int(v_i.shape[0]) if f_i.numel() == 0: continue @@ -1454,7 +1582,7 @@ class BakeNormalMapFromMesh(IO.ComfyNode): return IO.Schema( node_id="BakeNormalMapFromMesh", display_name="Bake Normal Map from Mesh", - category="latent/3d", + category="3d/mesh/texturing", description=( "Bakes a tangent-space normal map (glTF/OpenGL +Y) from a high-poly mesh into a " "low-poly's UV layout, capturing detail lost to decimation. Feed the UV-unwrapped " @@ -1495,7 +1623,7 @@ class BakeNormalMapFromMesh(IO.ComfyNode): imgs = [] for i in range(B): - v_i, f_i, _ = get_mesh_batch_item(low_poly, i) + v_i, f_i, *_ = get_mesh_batch_item(low_poly, i) n = int(v_i.shape[0]) if f_i.numel() == 0: logging.warning(f"BakeNormalMapFromMesh: skipping batch {i} (empty mesh)") @@ -1511,7 +1639,7 @@ class BakeNormalMapFromMesh(IO.ComfyNode): n_attr_i = low_n_attr[i, :n] if low_n_attr is not None else None low_n, tangents = _vertex_tangents_for_item(lv, lf, torch.from_numpy(uv_np).to(dev), n_attr_i, dev) - hv_i, hf_i, _ = get_mesh_batch_item(high_poly, i if h_batch > 1 else 0) + hv_i, hf_i, *_ = get_mesh_batch_item(high_poly, i if h_batch > 1 else 0) hv = hv_i.to(dev).float() hf = hf_i.to(dev).long() high_n = (high_n_attr[i, :hv.shape[0]].to(dev).float() if high_n_attr is not None @@ -1535,7 +1663,7 @@ class BakeAmbientOcclusion(IO.ComfyNode): return IO.Schema( node_id="BakeAmbientOcclusion", display_name="Bake Ambient Occlusion", - category="latent/3d", + category="3d/mesh/texturing", description=( "Bakes an ambient-occlusion map from a high-poly mesh into a low-poly's UV " "layout (white = open, dark = crevices). Feed the UV-unwrapped low_poly and the " @@ -1575,7 +1703,7 @@ class BakeAmbientOcclusion(IO.ComfyNode): pbar = comfy.utils.ProgressBar(max(1, B)) # one tick per batch item imgs = [] for i in range(B): - v_i, f_i, _ = get_mesh_batch_item(low_poly, i) + v_i, f_i, *_ = get_mesh_batch_item(low_poly, i) n = int(v_i.shape[0]) if f_i.numel() == 0: logging.warning(f"BakeAmbientOcclusion: skipping batch {i} (empty mesh)") @@ -1590,7 +1718,7 @@ class BakeAmbientOcclusion(IO.ComfyNode): low_n = (low_n_attr[i, :n].to(dev).float() if low_n_attr is not None else _compute_vertex_normals(lv, lf)) - hv_i, hf_i, _ = get_mesh_batch_item(high_poly, i if h_batch > 1 else 0) + hv_i, hf_i, *_ = get_mesh_batch_item(high_poly, i if h_batch > 1 else 0) img = _bake_ambient_occlusion( hv_i.to(dev).float(), hf_i.to(dev).long(), lv.detach().cpu().numpy(), lf.detach().cpu().numpy().astype(np.uint32), uv_np, @@ -1604,186 +1732,108 @@ class BakeAmbientOcclusion(IO.ComfyNode): return IO.NodeOutput(ao_img) -class SetMeshMaterial(IO.ComfyNode): +class RenderMesh(IO.ComfyNode): @classmethod def define_schema(cls): return IO.Schema( - node_id="SetMeshMaterial", - display_name="Set Mesh Material", - category="latent/3d", + node_id="RenderMesh", + display_name="Render Mesh", + search_aliases=["mesh to image", "render mesh", "preview textured mesh"], + category="3d/mesh", description=( - "Sets glTF material properties SaveGLB can't derive from textures: emissive " - "(color + strength + optional texture), baseColor tint, metallic/roughness " - "factors, normal scale, occlusion strength, double-sided. Place before SaveGLB." + "Ray-casts a single view of a mesh. The camera comes from a camera_info (Load3D / Preview3D viewer, or a Create Camera Info node)" ), inputs=[ IO.Mesh.Input("mesh"), - IO.Float.Input("emissive_r", default=0.0, min=0.0, max=1.0, step=0.01), - IO.Float.Input("emissive_g", default=0.0, min=0.0, max=1.0, step=0.01), - IO.Float.Input("emissive_b", default=0.0, min=0.0, max=1.0, step=0.01), - IO.Float.Input("emissive_strength", default=1.0, min=0.0, max=100.0, step=0.1, - tooltip=">1 for HDR glow (KHR_materials_emissive_strength)."), - IO.Image.Input("emissive_texture", optional=True), - IO.Float.Input("base_color_r", default=1.0, min=0.0, max=1.0, step=0.01), - IO.Float.Input("base_color_g", default=1.0, min=0.0, max=1.0, step=0.01), - IO.Float.Input("base_color_b", default=1.0, min=0.0, max=1.0, step=0.01), - IO.Float.Input("metallic_factor", default=-1.0, min=-1.0, max=1.0, step=0.01, - tooltip="-1 = leave auto; 0..1 overrides."), - IO.Float.Input("roughness_factor", default=-1.0, min=-1.0, max=1.0, step=0.01, - tooltip="-1 = leave auto; 0..1 overrides."), - IO.Float.Input("normal_scale", default=1.0, min=0.0, max=10.0, step=0.05), - IO.Float.Input("occlusion_strength", default=1.0, min=0.0, max=1.0, step=0.01), - IO.Boolean.Input("double_sided", default=True), + IO.Combo.Input("mode", options=["auto", "texture", "vertex colors", "solid", "normal", "depth"], + tooltip="What to render. auto: texture if present, else vertex colours, else shaded clay."), + IO.Int.Input("width", default=1024, min=64, max=4096, step=8), + IO.Int.Input("height", default=1024, min=64, max=4096, step=8), + IO.Color.Input("background", default="#000000"), + IO.Load3DCamera.Input("camera_info", optional=True, + tooltip="Camera from a Load3D / Preview3D viewer or a Create Camera Info " + "node. If none is connected, a default front view is auto-framed."), ], - outputs=[IO.Mesh.Output("mesh")], + outputs=[IO.Image.Output(display_name="image"), IO.Mask.Output(display_name="mask")], ) @classmethod - def execute(cls, mesh, emissive_r, emissive_g, emissive_b, emissive_strength, - base_color_r, base_color_g, base_color_b, metallic_factor, roughness_factor, - normal_scale, occlusion_strength, double_sided, emissive_texture=None): - out_mesh = copy.copy(mesh) - material = dict(mesh.material or {}) # merge over any prior material - material.update({ - "emissive_factor": [float(emissive_r), float(emissive_g), float(emissive_b)], - "emissive_strength": float(emissive_strength), - "base_color_factor": [float(base_color_r), float(base_color_g), float(base_color_b), 1.0], - "metallic_factor": float(metallic_factor), # <0 => leave auto - "roughness_factor": float(roughness_factor), - "normal_scale": float(normal_scale), - "occlusion_strength": float(occlusion_strength), - "double_sided": bool(double_sided), - }) - out_mesh.material = material - if emissive_texture is not None: - out_mesh.emissive = emissive_texture.float().clamp(0.0, 1.0).cpu() - return IO.NodeOutput(out_mesh) + def execute(cls, mesh, mode, width, height, background, camera_info=None): + if int(mesh.vertices.shape[0]) > 1: + logging.warning("RenderMesh: one item per batch only; using the first of %d.", int(mesh.vertices.shape[0])) + dev = comfy.model_management.get_torch_device() + v_i, f_i, *_ = get_mesh_batch_item(mesh, 0) + n = int(v_i.shape[0]) + lv = v_i.to(dev).float() + lf = f_i.to(dev).long() + def _item(attr): # first-item, length-n view of a mesh attr + a = getattr(mesh, attr, None) + if a is None: + return None + return (a[0] if (isinstance(a, list) or a.ndim == 3) else a)[:n].to(dev).float() -def paint_mesh_default_colors(mesh): - out_mesh = copy.copy(mesh) - vertex_count = mesh.vertices.shape[1] - out_mesh.vertex_colors = mesh.vertices.new_zeros((1, vertex_count, 3)) - return out_mesh + tex, uvs = mesh.texture, mesh.uvs + have_tex = tex is not None and uvs is not None + have_vc = getattr(mesh, "vertex_colors", None) is not None + resolved = mode + if resolved == "auto": + resolved = "texture" if have_tex else ("vertex colors" if have_vc else "solid") + uv = texture_hwc = vcols = vnorms = None + render_normal = False + if resolved == "texture" and have_tex: + uv = (uvs[0, :n] if uvs.ndim == 3 else uvs[:n]).to(dev).float() + texture_hwc = tex[0].to(dev).float() + elif resolved == "vertex colors" and have_vc: + # glTF COLOR_0 is linear (PaintMesh stores pow(srgb, 2.2)); to sRGB so it isn't dark. + vcols = _item("vertex_colors")[:, :3].clamp(0.0, 1.0).pow(1.0 / 2.2) + elif resolved == "normal": + render_normal = True + vnorms = _item("normals") # smooth if present, else face normals -def fill_holes_fn(vertices, faces, max_perimeter=0.03): - is_batched = vertices.ndim == 3 - if is_batched: - v_list, f_list = [], [] - for i in range(vertices.shape[0]): - v_i, f_i = fill_holes_fn(vertices[i], faces[i], max_perimeter) - v_list.append(v_i) - f_list.append(f_i) - max_v = max(v.shape[0] for v in v_list) - for i in range(len(v_list)): - if v_list[i].shape[0] < max_v: - pad = torch.zeros(max_v - v_list[i].shape[0], 3, device=v_list[i].device, dtype=v_list[i].dtype) - v_list[i] = torch.cat([v_list[i], pad], dim=0) - return torch.stack(v_list), torch.stack(f_list) + tri = lv[lf] + bvh = _build_triangle_bvh(tri) + center = (lv.amax(0) + lv.amin(0)) * 0.5 + up_hint = torch.tensor([0.0, 1.0, 0.0], device=dev) - device = vertices.device - v = vertices - f = faces - - if f.numel() == 0: - return v, f - - edges = torch.cat([f[:, [0, 1]], f[:, [1, 2]], f[:, [2, 0]]], dim=0) - edges_sorted, _ = torch.sort(edges, dim=1) - max_v = v.shape[0] - packed_undirected = edges_sorted[:, 0].long() * max_v + edges_sorted[:, 1].long() - unique_packed, counts = torch.unique(packed_undirected, return_counts=True) - boundary_packed = unique_packed[counts == 1] - - if boundary_packed.numel() == 0: - return v, f - - boundary_mask = torch.isin(packed_undirected, boundary_packed) - b_edges = edges_sorted[boundary_mask] - - adj = {} - for i in range(b_edges.shape[0]): - a = b_edges[i, 0].item() - b = b_edges[i, 1].item() - adj.setdefault(a, []).append(b) - adj.setdefault(b, []).append(a) - - # Trace all boundary loops - loops = [] - visited = set() - for start_node in adj.keys(): - if start_node in visited: - continue - curr = start_node - prev = -1 - loop = [] - while curr not in visited: - visited.add(curr) - loop.append(curr) - neighbors = adj[curr] - candidates = [n for n in neighbors if n != prev] - if not candidates: - loop = [] - break - next_node = candidates[0] - prev, curr = curr, next_node - if curr == start_node: - loops.append(loop) - break - - if not loops: - return v, f - - # Mesh normal for winding orientation only - face_normals = torch.linalg.cross( - v[f[:, 1]] - v[f[:, 0]], - v[f[:, 2]] - v[f[:, 0]], - dim=-1 - ) - mesh_normal = face_normals.mean(dim=0) - mesh_normal = mesh_normal / (torch.norm(mesh_normal) + 1e-8) - - # Fill all boundary loops below the perimeter threshold. - new_verts = [] - new_faces = [] - v_idx = v.shape[0] - - for loop in loops: - loop_t = torch.tensor(loop, device=device, dtype=torch.long) - loop_v = v[loop_t] - - next_v = torch.roll(loop_v, -1, dims=0) - diffs = loop_v - next_v - perimeter = torch.norm(diffs, dim=1).sum().item() - - if perimeter > max_perimeter: - continue - - # Ensure CCW winding consistent with mesh - cross = torch.linalg.cross(loop_v, next_v, dim=-1) - loop_normal = cross.sum(dim=0) - loop_normal = loop_normal / (torch.norm(loop_normal) + 1e-8) - if torch.dot(loop_normal, mesh_normal) < 0: - loop = loop[::-1] - loop_t = torch.tensor(loop, device=device, dtype=torch.long) - loop_v = v[loop_t] - - if len(loop) == 3: - new_faces.append([loop[0], loop[1], loop[2]]) + if camera_info is not None: + # Explicit camera from a Load3D / Preview3D viewer (three.js Y-up, same frame as the mesh). + def _vec(d): + return torch.tensor([float(d.get("x", 0.0)), float(d.get("y", 0.0)), float(d.get("z", 0.0))], + device=dev) + eye = _vec(camera_info.get("position", {})) + tgt = camera_info.get("target") + target = _vec(tgt) if tgt else center + f, r, u = _camera_basis(eye, target, up_hint) # look-at (roll-free) + cam_fov = float(camera_info.get("fov", 0) or 0) or 40.0 + cam_zoom = float(camera_info.get("zoom", 1.0) or 1.0) # three.js digital zoom scales focal length + fov_rad = 2.0 * math.atan(math.tan(math.radians(cam_fov) * 0.5) / max(cam_zoom, 1e-3)) else: - centroid = loop_v.mean(dim=0) - new_verts.append(centroid) - for i in range(len(loop)): - new_faces.append([loop[i], loop[(i + 1) % len(loop)], v_idx]) - v_idx += 1 + # No camera connected, auto-framed default front view. + fov_rad = math.radians(40.0) + r_sphere = float((lv - center).norm(dim=-1).amax().clamp_min(1e-6)) + radius = r_sphere / math.tan(fov_rad * 0.5) * 1.04 + eye = center + torch.tensor([0.0, 0.0, 1.0], device=dev) * radius + f, r, u = _camera_basis(eye, center, up_hint) - if new_verts: - v = torch.cat([v, torch.stack(new_verts)], dim=0) - if new_faces: - f = torch.cat([f, torch.tensor(new_faces, device=device, dtype=torch.long)], dim=0) + H, W = int(height), int(width) + img, hit, depth = _render_view(tri, bvh, uv, lf, texture_hwc, eye, f, r, u, fov_rad, H, W, + vertex_colors=vcols, vertex_normals=vnorms, render_normal=render_normal) - return v, f + if resolved == "depth": + img = torch.zeros_like(img) + if bool(hit.any()): + dh = depth[hit] + rng = max(float(dh.max()) - float(dh.min()), 1e-6) + norm = ((float(dh.max()) - depth) / rng).clamp(0.0, 1.0) # near (small depth) = white + img = torch.where(hit[..., None], norm[..., None].expand(-1, -1, 3), img) + + bg = background.lstrip("#") + bg_rgb = torch.tensor([int(bg[i:i + 2], 16) / 255.0 for i in (0, 2, 4)], device=dev) + out = torch.where(hit[..., None], img.clamp(0.0, 1.0), bg_rgb) + idev, idtype = comfy.model_management.intermediate_device(), comfy.model_management.intermediate_dtype() + return IO.NodeOutput(out[None].to(idev, idtype), hit.float()[None].to(idev, idtype)) def _fill_holes_v2_gpu(verts, faces, max_perimeter, colors=None, fill_chains=False, max_verts=16): @@ -2008,8 +2058,7 @@ def weld_vertices_fn(vertices, faces, epsilon=None, epsilon_rel=1e-5, colors=Non def fill_holes_v2_fn(vertices, faces, max_perimeter=0.03, colors=None, weld_epsilon_rel=1e-5, fill_chains=False, max_verts=16): - """Batched v2 GPU hole-filler (v1 CPU walker fallback on non-CUDA). Pre-welds verts - first — boundary detection needs shared edges; `weld_epsilon_rel=0` skips it.""" + """Batched v2 GPU hole-filler. Pre-welds verts first as boundary detection needs shared edges""" if vertices.ndim == 3: v_list, f_list, c_list = [], [], [] if colors is not None else None pbar = comfy.utils.ProgressBar(vertices.shape[0]) @@ -2063,12 +2112,11 @@ def fill_holes_v2_fn(vertices, faces, max_perimeter=0.03, colors=None, weld_epsi f"duplicate verts at distances >{WELD_CAP}× bbox; fix upstream " f"(decimate node settings) or run WeldVertices manually with a larger epsilon." ) - if vertices.device.type == "cuda": - out_v, out_f, out_c, _ = _fill_holes_v2_gpu(vertices, faces, max_perimeter, colors, fill_chains, max_verts) - return out_v, out_f, out_c - # CPU fallback: v1 walker (no attribute prop, but topologically equivalent for manifold boundary). - out_v, out_f = fill_holes_fn(vertices, faces, max_perimeter=max_perimeter) - return out_v, out_f, colors + dev = comfy.model_management.get_torch_device() + out_v, out_f, out_c, _ = _fill_holes_v2_gpu( + vertices.to(dev), faces.to(dev), max_perimeter, + colors.to(dev) if colors is not None else None, fill_chains, max_verts) + return out_v, out_f, out_c def _process_mesh_batch(mesh, per_item_fn): @@ -2161,7 +2209,7 @@ class DecimateMesh(IO.ComfyNode): return IO.Schema( node_id="DecimateMesh", display_name="Decimate Mesh", - category="latent/3d", + category="3d/mesh", description=( "Simplifies a mesh to a target face count using QEM, on the active compute " "device. 'midpoint' is the cumesh-faithful preset (best quality, preserves thin " @@ -2252,7 +2300,7 @@ class RemeshMesh(IO.ComfyNode): return IO.Schema( node_id="RemeshMesh", display_name="Remesh Mesh (Narrow-Band DC)", - category="latent/3d", + category="3d/mesh", description=( "Re-extracts a uniformly tessellated mesh via a narrow-band distance field + Dual " "Contouring, on the active compute device. Normalizes messy / non-manifold / " @@ -2422,8 +2470,6 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance "(triangle soup); UV charts will be per-face. Raise weld_distance.") if segmenter == "pec": - if mesh.faces.device.type != "cuda": - raise RuntimeError("segmenter='pec' requires a CUDA mesh; use 'adaptive' for CPU.") face_chart = _uv_seg.cluster_charts_pec(mesh, max_cost=1.0) elif segmenter == "adaptive": face_chart = _uv_seg.segment_charts(mesh, max_cost=2.0) @@ -2508,7 +2554,7 @@ class UnwrapMesh(IO.ComfyNode): return IO.Schema( node_id="UnwrapMesh", display_name="Unwrap Mesh UVs", - category="latent/3d", + category="3d/mesh/texturing", description=( "Generates a UV atlas (pure-torch, no xatlas): segments the surface into charts, " "parameterizes each, packs into a [0,1] atlas. Seam verts duplicated. Run after " @@ -2517,8 +2563,7 @@ class UnwrapMesh(IO.ComfyNode): inputs=[ IO.Mesh.Input("mesh"), IO.Combo.Input("segmenter", options=["pec", "adaptive"], default="pec", - tooltip="pec: fast parallel-edge-collapse charting (CUDA; CPU falls back to " - "adaptive). adaptive: CPU, slower."), + tooltip="pec: fast parallel-edge-collapse charting on GPU. adaptive: CPU, slower."), IO.Int.Input("resolution", default=1024, min=0, max=8192, step=256, tooltip="Target atlas resolution for texel-density auto-scale (0 = fit-to-content)."), IO.Int.Input("padding", default=1, min=0, max=16, @@ -2534,10 +2579,7 @@ class UnwrapMesh(IO.ComfyNode): @classmethod def execute(cls, mesh, segmenter, resolution, padding, weld_distance): compute_device = comfy.model_management.get_torch_device() - seg = segmenter - if seg == "pec" and compute_device.type != "cuda": - seg = "adaptive" - seg_device = compute_device if seg == "pec" else torch.device("cpu") + seg_device = compute_device if segmenter == "pec" else torch.device("cpu") is_list = isinstance(mesh.vertices, list) is_batched = not is_list and mesh.vertices.ndim == 3 @@ -2563,7 +2605,7 @@ class UnwrapMesh(IO.ComfyNode): vmapping, indices, uvs = _uv_unwrap( vi.to(seg_device).float(), fi.to(seg_device).long(), - seg, int(resolution), int(padding), weld_abs) + segmenter, int(resolution), int(padding), weld_abs) uvs = uvs.copy() uvs[:, 1] = 1.0 - uvs[:, 1] # UV y flipped vs trimesh @@ -2746,9 +2788,8 @@ class RenderUVAtlas(IO.ComfyNode): return IO.Schema( node_id="RenderUVAtlas", display_name="Render UV Atlas", - category="latent/3d", - description=("Renders a mesh's UV layout as an image (each chart a distinct color, " - "outlined at borders). Run UnwrapMesh first."), + category="3d/mesh/texturing", + description=("Renders a mesh's UV layout as an image."), inputs=[ IO.Mesh.Input("mesh"), IO.Int.Input("resolution", default=1024, min=64, max=4096, step=64), @@ -2782,11 +2823,9 @@ class FillHoles(IO.ComfyNode): return IO.Schema( node_id="FillHoles", display_name="Fill Holes", - category="latent/3d", + category="3d/mesh", description=( - "Fills holes up to a max perimeter, preserving existing geometry/UVs (only patch " - "tris added). GPU-vectorised with auto-corrected winding and loop-averaged centroid " - "colors; CPU walker fallback on non-CUDA." + "Fills holes up to a max perimeter, preserving existing geometry/UVs. GPU-vectorised with auto-corrected winding and loop-averaged centroid colors" ), inputs=[ IO.Mesh.Input("mesh"), @@ -2824,7 +2863,7 @@ class WeldVertices(IO.ComfyNode): return IO.Schema( node_id="WeldVertices", display_name="Weld Vertices", - category="latent/3d", + category="3d/mesh", description=( "Merge coincident vertices via L_inf grid quantization. Use when a mesh comes in " "unwelded (per-face verts, no shared edges) — pre-pass before FillHoles, " @@ -2852,83 +2891,85 @@ class WeldVertices(IO.ComfyNode): return _process_mesh_batch(mesh, _fn) -def merge_meshes(meshes): - """Concatenate Types.MESH list into one (B=1) mesh: cumulative face-index offset, - missing uvs/colors padded (zeros/white), texture from the first input that has one - (later dropped — single-primitive glb can't carry multiple atlases).""" - if not meshes: - raise ValueError("merge_meshes: need at least one mesh") - - def _b0(t): - return t[0] if t.ndim == 3 else t - - any_uvs = any(m.uvs is not None for m in meshes) - any_colors = any(m.vertex_colors is not None for m in meshes) - - verts_list, faces_list, uvs_list, colors_list = [], [], [], [] - texture = None - offset = 0 - for m in meshes: - # Coerce to CPU so CUDA-side (MoGe) meshes merge cleanly with our outputs. - v = _b0(m.vertices).cpu() - f = _b0(m.faces).cpu() - verts_list.append(v) - faces_list.append(f + offset) - offset += v.shape[0] - if any_uvs: - mu = m.uvs - uvs_list.append(_b0(mu).cpu() if mu is not None else v.new_zeros((v.shape[0], 2))) - if any_colors: - mc = m.vertex_colors - if mc is not None: - c = _b0(mc).cpu() - else: - c = v.new_ones((v.shape[0], 3)) - colors_list.append(c) - mt = m.texture - if mt is not None: - if texture is None: - texture = mt.cpu() - else: - logging.warning("MergeMeshes: dropping extra texture from input; only one texture is kept.") - - merged_verts = torch.cat(verts_list, dim=0).unsqueeze(0) - merged_faces = torch.cat(faces_list, dim=0).unsqueeze(0) - merged_uvs = torch.cat(uvs_list, dim=0).unsqueeze(0) if any_uvs else None - merged_colors = torch.cat(colors_list, dim=0).unsqueeze(0) if any_colors else None - - return Types.MESH( - vertices=merged_verts, - faces=merged_faces, - uvs=merged_uvs, - vertex_colors=merged_colors, - texture=texture, - ) - - -class MergeMeshes(IO.ComfyNode): +class MeshSmoothNormals(IO.ComfyNode): @classmethod def define_schema(cls): - autogrow_template = IO.Autogrow.TemplatePrefix( - IO.Mesh.Input("mesh"), prefix="mesh", min=2, max=50, - ) return IO.Schema( - node_id="MergeMeshes", - display_name="Merge Meshes", - category="latent/3d", + node_id="MeshSmoothNormals", + display_name="Smooth Mesh Normals", + category="3d", description=( - "Concatenate N meshes into one by offsetting face indices and stacking verts, " - "faces, uvs, and colors." + "Compute smooth per-vertex normals and attach them to the mesh. Meshes " + "without normals are shaded flat (per-face) by glTF viewers; this makes " + "them shade smoothly. With crease_angle below 180, edges sharper than the " + "threshold are kept hard by splitting vertices along them." ), inputs=[ - IO.Autogrow.Input("meshes", template=autogrow_template), + IO.Mesh.Input("mesh"), + IO.Float.Input("crease_angle", default=180.0, min=0.0, max=180.0, step=1.0, + tooltip="Edges whose dihedral angle exceeds this (degrees) stay " + "hard (vertices are split). 180 = fully smooth; lower " + "preserves sharp edges (e.g. ~30-60 for hard-surface)."), ], outputs=[IO.Mesh.Output("mesh")], ) @classmethod - def execute(cls, meshes: IO.Autogrow.Type) -> IO.NodeOutput: - return IO.NodeOutput(merge_meshes(list(meshes.values()))) + def execute(cls, mesh: Types.MESH, crease_angle: float) -> IO.NodeOutput: + crease = None if crease_angle >= 180.0 else float(crease_angle) + batch_size = mesh.vertices.shape[0] + + if crease is None: + # Fully smooth: topology is unchanged, so just attach a normals tensor that + # matches the existing (possibly zero-padded) vertex layout and keep all fields. + normals_padded = torch.zeros_like(mesh.vertices) + for i in range(batch_size): + v_i, f_i, _, _, _ = get_mesh_batch_item(mesh, i) + if v_i.shape[0] == 0 or f_i.shape[0] == 0: + continue + n_i = _smooth_vertex_normals(v_i.cpu().numpy().astype(np.float32), + f_i.cpu().numpy().astype(np.int64)) + normals_padded[i, :n_i.shape[0]] = torch.from_numpy(n_i).to(mesh.vertices) + out = copy.copy(mesh) + out.normals = normals_padded + return IO.NodeOutput(out) + + # Crease split changes per-item vertex counts -> rebuild as a variable-size batch. + tangents_b = mesh.tangents + v_list, f_list, n_list = [], [], [] + c_list = [] if mesh.vertex_colors is not None else None + u_list = [] if mesh.uvs is not None else None + t_list = [] if tangents_b is not None else None + for i in range(batch_size): + v_i, f_i, c_i, u_i, _ = get_mesh_batch_item(mesh, i) + if v_i.shape[0] == 0 or f_i.shape[0] == 0: + continue + dev = v_i.device + vo, fo, no, remap = _compute_vertex_face_normals( + v_i.cpu().numpy().astype(np.float32), + f_i.cpu().numpy().astype(np.int64), crease) + remap_t = torch.from_numpy(remap) + v_list.append(torch.from_numpy(vo).to(dev, mesh.vertices.dtype)) + f_list.append(torch.from_numpy(fo.astype(np.int64)).to(dev, mesh.faces.dtype)) + n_list.append(torch.from_numpy(no).to(dev, mesh.vertices.dtype)) + if c_list is not None: + c_list.append(c_i[remap_t.to(c_i.device)]) + if u_list is not None: + u_list.append(u_i[remap_t.to(u_i.device)]) + if t_list is not None: + # Remap (not recompute) so TANGENT keeps the baked basis; split verts copy theirs. + t_i = tangents_b[i, :v_i.shape[0]] + t_list.append(t_i[remap_t.to(t_i.device)]) + if not v_list: + return IO.NodeOutput(mesh) + out = pack_variable_mesh_batch( + v_list, f_list, colors=c_list, uvs=u_list, + texture=mesh.texture, unlit=mesh.unlit, + normals=n_list, metallic_roughness=mesh.metallic_roughness, + tangents=t_list, normal_map=mesh.normal_map, + occlusion_in_mr=mesh.occlusion_in_mr, + material=mesh.material, emissive=mesh.emissive) + return IO.NodeOutput(out) class PostProcessMeshExtension(ComfyExtension): @@ -2947,8 +2988,8 @@ class PostProcessMeshExtension(ComfyExtension): ApplyTextureToMesh, BakeNormalMapFromMesh, BakeAmbientOcclusion, - SetMeshMaterial, - MergeMeshes, + RenderMesh, + MeshSmoothNormals, ] diff --git a/comfy_extras/nodes_save_3d.py b/comfy_extras/nodes_save_3d.py index 07e0befd2..6b6eec51d 100644 --- a/comfy_extras/nodes_save_3d.py +++ b/comfy_extras/nodes_save_3d.py @@ -105,104 +105,6 @@ def get_mesh_batch_item(mesh, index): return mesh.vertices[index], mesh.faces[index], colors, uvs, normals -def _smooth_vertex_normals(vertices_np, faces_np, weld=True): - """Area-weighted per-vertex normals (unit length), fully smooth — no vertex splitting. - - Un-normalized face normals (the raw cross product) have magnitude 2*area, so - accumulating them onto their vertices yields an area-weighted average. `weld` averages - across vertices that share a position — UV-seam duplicates created by unwrapping — so - both sides of a seam get one identical normal. Without it each side averages only its - own faces and a visible shading seam appears; welding matches the official, which - computes normals on the pre-split mesh and gathers them through the UV vmap.""" - tris = vertices_np[faces_np] # (M, 3, 3) - face_n = np.cross(tris[:, 1] - tris[:, 0], tris[:, 2] - tris[:, 0]) - if weld and vertices_np.shape[0]: - # Group coincident positions (quantized to ~1e-5 of the bbox) into one shared normal. - lo = vertices_np.min(0) - inv_tol = 1.0 / (max(float((vertices_np.max(0) - lo).max()), 1e-9) * 1e-5) - q = np.round((vertices_np - lo) * inv_tol).astype(np.int64) - _, group = np.unique(q, axis=0, return_inverse=True) - acc = np.zeros((int(group.max()) + 1, 3), dtype=np.float64) - for k in range(3): - np.add.at(acc, group[faces_np[:, k]], face_n) - normals = acc[group] # welded normal back to each vertex - else: - normals = np.zeros((vertices_np.shape[0], 3), dtype=np.float64) - for k in range(3): - np.add.at(normals, faces_np[:, k], face_n) - lens = np.linalg.norm(normals, axis=1, keepdims=True) - normals /= np.where(lens > 1e-12, lens, 1.0) - return normals.astype(np.float32) - - -def _compute_vertex_normals(vertices_np, faces_np, crease_angle=None): - """Compute per-vertex normals, returning (vertices, faces_uint32, normals, remap). - - crease_angle is None (or >= 180) -> fully smooth normals; vertices/faces are - returned unchanged and remap is None. - - Otherwise vertices are split along edges whose dihedral angle exceeds - crease_angle (degrees) so hard creases stay sharp while smooth regions still - interpolate. remap maps each output vertex back to its source index, so the - caller can duplicate any per-vertex attributes (uvs / colors) to match.""" - faces_i = faces_np.astype(np.int64) - if crease_angle is None or crease_angle >= 180.0: - return (vertices_np, faces_i.astype(np.uint32), - _smooth_vertex_normals(vertices_np, faces_i), None) - - M = faces_i.shape[0] - tris = vertices_np[faces_i] - face_n = np.cross(tris[:, 1] - tris[:, 0], tris[:, 2] - tris[:, 0]) - areas = np.linalg.norm(face_n, axis=1, keepdims=True) - face_unit = face_n / np.where(areas > 1e-12, areas, 1.0) - cos_thresh = math.cos(math.radians(crease_angle)) - - # Union faces that share an edge whose dihedral angle is below the crease - # threshold; each connected component becomes one smoothing group. - parent = list(range(M)) - - def find(x): - while parent[x] != x: - parent[x] = parent[parent[x]] - x = parent[x] - return x - - edge_faces = {} - for fi in range(M): - a, b, c = int(faces_i[fi, 0]), int(faces_i[fi, 1]), int(faces_i[fi, 2]) - for u, v in ((a, b), (b, c), (c, a)): - edge_faces.setdefault((u, v) if u < v else (v, u), []).append(fi) - for fl in edge_faces.values(): - if len(fl) == 2 and float(np.dot(face_unit[fl[0]], face_unit[fl[1]])) >= cos_thresh: - ra, rb = find(fl[0]), find(fl[1]) - if ra != rb: - parent[ra] = rb - - # Emit one output vertex per (original vertex, smoothing group) pair. - new_index = {} - remap = [] - out_faces = np.empty((M, 3), dtype=np.int64) - for fi in range(M): - g = find(fi) - for k in range(3): - ov = int(faces_i[fi, k]) - key = (ov, g) - ni = new_index.get(key) - if ni is None: - ni = len(remap) - new_index[key] = ni - remap.append(ov) - out_faces[fi, k] = ni - - remap = np.asarray(remap, dtype=np.int64) - normals = np.zeros((remap.shape[0], 3), dtype=np.float64) - for k in range(3): - np.add.at(normals, out_faces[:, k], face_n) - lens = np.linalg.norm(normals, axis=1, keepdims=True) - normals /= np.where(lens > 1e-12, lens, 1.0) - return (vertices_np[remap], out_faces.astype(np.uint32), normals.astype(np.float32), remap) - - def save_glb(vertices, faces, filepath=None, metadata=None, uvs=None, vertex_colors=None, texture_image=None, metallic_roughness_image=None, unlit=False, @@ -823,91 +725,83 @@ class RotateMesh(IO.ComfyNode): return IO.NodeOutput(out) -class MeshSmoothNormals(IO.ComfyNode): +class MergeMeshes(IO.ComfyNode): @classmethod def define_schema(cls): + autogrow_template = IO.Autogrow.TemplatePrefix( + IO.Mesh.Input("mesh"), prefix="mesh", min=2, max=50, + ) return IO.Schema( - node_id="MeshSmoothNormals", - display_name="Smooth Mesh Normals", - category="3d", + node_id="MergeMeshes", + display_name="Merge Meshes", + category="3d/mesh", description=( - "Compute smooth per-vertex normals and attach them to the mesh. Meshes " - "without normals are shaded flat (per-face) by glTF viewers; this makes " - "them shade smoothly. With crease_angle below 180, edges sharper than the " - "threshold are kept hard by splitting vertices along them." + "Concatenate N meshes into one by offsetting face indices and stacking verts, " + "faces, uvs, and colors." ), inputs=[ - IO.Mesh.Input("mesh"), - IO.Float.Input("crease_angle", default=180.0, min=0.0, max=180.0, step=1.0, - tooltip="Edges whose dihedral angle exceeds this (degrees) stay " - "hard (vertices are split). 180 = fully smooth; lower " - "preserves sharp edges (e.g. ~30-60 for hard-surface)."), + IO.Autogrow.Input("meshes", template=autogrow_template), ], outputs=[IO.Mesh.Output("mesh")], ) @classmethod - def execute(cls, mesh: Types.MESH, crease_angle: float) -> IO.NodeOutput: - crease = None if crease_angle >= 180.0 else float(crease_angle) - batch_size = mesh.vertices.shape[0] + def execute(cls, meshes: IO.Autogrow.Type) -> IO.NodeOutput: + # Concatenate the input meshes into one (B=1) mesh: cumulative face-index offset, + # missing uvs/colors padded (zeros/white), texture from the first input that has one + # (later dropped — a single-primitive glb can't carry multiple atlases). + meshes = list(meshes.values()) + if not meshes: + raise ValueError("MergeMeshes: need at least one mesh") - if crease is None: - # Fully smooth: topology is unchanged, so just attach a normals tensor that - # matches the existing (possibly zero-padded) vertex layout and keep all fields. - normals_padded = torch.zeros_like(mesh.vertices) - for i in range(batch_size): - v_i, f_i, _, _, _ = get_mesh_batch_item(mesh, i) - if v_i.shape[0] == 0 or f_i.shape[0] == 0: - continue - n_i = _smooth_vertex_normals(v_i.cpu().numpy().astype(np.float32), - f_i.cpu().numpy().astype(np.int64)) - normals_padded[i, :n_i.shape[0]] = torch.from_numpy(n_i).to(mesh.vertices) - out = copy.copy(mesh) - out.normals = normals_padded - return IO.NodeOutput(out) + def _b0(t): + return t[0] if t.ndim == 3 else t - # Crease split changes per-item vertex counts -> rebuild as a variable-size batch. - tangents_b = mesh.tangents - v_list, f_list, n_list = [], [], [] - c_list = [] if mesh.vertex_colors is not None else None - u_list = [] if mesh.uvs is not None else None - t_list = [] if tangents_b is not None else None - for i in range(batch_size): - v_i, f_i, c_i, u_i, _ = get_mesh_batch_item(mesh, i) - if v_i.shape[0] == 0 or f_i.shape[0] == 0: - continue - dev = v_i.device - vo, fo, no, remap = _compute_vertex_normals( - v_i.cpu().numpy().astype(np.float32), - f_i.cpu().numpy().astype(np.int64), crease) - remap_t = torch.from_numpy(remap) - v_list.append(torch.from_numpy(vo).to(dev, mesh.vertices.dtype)) - f_list.append(torch.from_numpy(fo.astype(np.int64)).to(dev, mesh.faces.dtype)) - n_list.append(torch.from_numpy(no).to(dev, mesh.vertices.dtype)) - if c_list is not None: - c_list.append(c_i[remap_t.to(c_i.device)]) - if u_list is not None: - u_list.append(u_i[remap_t.to(u_i.device)]) - if t_list is not None: - # Remap (not recompute) so TANGENT keeps the baked basis; split verts copy theirs. - t_i = tangents_b[i, :v_i.shape[0]] - t_list.append(t_i[remap_t.to(t_i.device)]) - if not v_list: - return IO.NodeOutput(mesh) - out = pack_variable_mesh_batch( - v_list, f_list, colors=c_list, uvs=u_list, - texture=mesh.texture, unlit=mesh.unlit, - normals=n_list, metallic_roughness=mesh.metallic_roughness, - tangents=t_list, normal_map=mesh.normal_map, - occlusion_in_mr=mesh.occlusion_in_mr, - material=mesh.material, emissive=mesh.emissive) - return IO.NodeOutput(out) + any_uvs = any(m.uvs is not None for m in meshes) + any_colors = any(m.vertex_colors is not None for m in meshes) + + verts_list, faces_list, uvs_list, colors_list = [], [], [], [] + texture = None + offset = 0 + for m in meshes: + # Coerce to CPU so CUDA-side (MoGe) meshes merge cleanly with our outputs. + v = _b0(m.vertices).cpu() + f = _b0(m.faces).cpu() + verts_list.append(v) + faces_list.append(f + offset) + offset += v.shape[0] + if any_uvs: + mu = m.uvs + uvs_list.append(_b0(mu).cpu() if mu is not None else v.new_zeros((v.shape[0], 2))) + if any_colors: + mc = m.vertex_colors + c = _b0(mc).cpu() if mc is not None else v.new_ones((v.shape[0], 3)) + colors_list.append(c) + mt = m.texture + if mt is not None: + if texture is None: + texture = mt.cpu() + else: + logging.warning("MergeMeshes: dropping extra texture from input; only one texture is kept.") + + merged_verts = torch.cat(verts_list, dim=0).unsqueeze(0) + merged_faces = torch.cat(faces_list, dim=0).unsqueeze(0) + merged_uvs = torch.cat(uvs_list, dim=0).unsqueeze(0) if any_uvs else None + merged_colors = torch.cat(colors_list, dim=0).unsqueeze(0) if any_colors else None + + return IO.NodeOutput(Types.MESH( + vertices=merged_verts, + faces=merged_faces, + uvs=merged_uvs, + vertex_colors=merged_colors, + texture=texture, + )) class Save3DExtension(ComfyExtension): @override async def get_node_list(self) -> list[type[IO.ComfyNode]]: - return [SaveGLB, MeshToFile3D, RotateMesh, MeshSmoothNormals] + return [SaveGLB, MeshToFile3D, RotateMesh, MergeMeshes] async def comfy_entrypoint() -> Save3DExtension: