From 0eb098b5910f60d3a73f4f027822cd3599396674 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:01:51 -0700 Subject: [PATCH] Lower peak trellis workflow vram and ram usage. (#16054) --- comfy/ldm/trellis2/flexgemm.py | 52 +++-- comfy/ldm/trellis2/vae.py | 74 +++--- comfy_api/latest/_util/geometry_types.py | 23 ++ comfy_execution/caching.py | 17 +- comfy_extras/mesh3d/uv_unwrap/pack.py | 272 +++++++++++++---------- comfy_extras/nodes_mesh_postprocess.py | 149 +++++++++---- comfy_extras/nodes_trellis2.py | 33 ++- 7 files changed, 397 insertions(+), 223 deletions(-) diff --git a/comfy/ldm/trellis2/flexgemm.py b/comfy/ldm/trellis2/flexgemm.py index 8fb6d56f7..76023992e 100644 --- a/comfy/ldm/trellis2/flexgemm.py +++ b/comfy/ldm/trellis2/flexgemm.py @@ -19,7 +19,7 @@ class TorchHashMap: def __init__(self, keys: torch.Tensor, values: torch.Tensor): self.sorted_keys, order = torch.sort(keys.to(torch.long)) - self.sorted_vals = values.to(torch.long)[order] + self.sorted_vals = values[order] self._n = self.sorted_keys.numel() # Chunk size for lookup_flat, caps each transient to ~CHUNK rows. @@ -54,9 +54,9 @@ def build_submanifold_neighbor_map( # Chunked over voxels so the [chunk, V, 3] candidate transient stays bounded. device = coords.device M = coords.shape[0] - offsets = compute_kernel_offsets(Kw, Kh, Kd, Dw, Dh, Dd, device).long() # [V, 3] + offsets = compute_kernel_offsets(Kw, Kh, Kd, Dw, Dh, Dd, device) # [V, 3] V = offsets.shape[0] - center = torch.tensor([(Kw // 2) * Dw, (Kh // 2) * Dh, (Kd // 2) * Dd], device=device) + center = torch.tensor([(Kw // 2) * Dw, (Kh // 2) * Dh, (Kd // 2) * Dd], dtype=torch.int32, device=device) WHD, HD = W * H * D, H * D neighbor = torch.empty((M, V), dtype=torch.int32, device=device) @@ -66,11 +66,12 @@ def build_submanifold_neighbor_map( for s in range(0, M, chunk): e = min(s + chunk, M) b = coords[s:e, 0].long() - cand = coords[s:e, 1:4].long()[:, None, :] + offsets[None, :, :] - center # [c, V, 3] + cand = coords[s:e, 1:4][:, None, :] + offsets[None, :, :] - center # [c, V, 3] x, y, z = cand[..., 0], cand[..., 1], cand[..., 2] in_bounds = (x >= 0) & (x < W) & (y >= 0) & (y < H) & (z >= 0) & (z < D) # [c, V] - flat = b[:, None] * WHD + x * HD + y * D + z # [c, V] - flat = torch.where(in_bounds, flat, torch.full_like(flat, -1)) # OOB -> guaranteed miss + flat = x.long().mul_(HD) + flat.add_(y.long().mul_(D)).add_(z).add_(b[:, None] * WHD) + flat.masked_fill_(~in_bounds, -1) # OOB -> guaranteed miss neighbor[s:e] = hashmap.lookup_flat(flat.reshape(-1)).view(e - s, V) return neighbor @@ -78,7 +79,7 @@ def get_recommended_chunk_mem( device=None, safety_fraction: float = 0.2, min_gb: float = 0.25, - max_gb: float = 2.0, + max_gb: float = 0.5, ): """Pick a chunk-memory budget (in GB) for sparse conv batching.""" free_gb = comfy.model_management.get_free_memory(device) / (1024 ** 3) @@ -92,6 +93,7 @@ def sparse_submanifold_conv3d( bias: Optional[torch.Tensor], neighbor_cache: Optional[torch.Tensor], dilation: tuple, + cache_neighbor_map: bool = True, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: if feats.shape[0] == 0: Co = weight.shape[0] @@ -103,23 +105,27 @@ def sparse_submanifold_conv3d( V = Kw * Kh * Kd device = feats.device + hashmap = None if neighbor_cache is None: b_stride = W * H * D x_stride = H * D y_stride = D z_stride = 1 - flat_keys = (coords[:, 0].long() * b_stride + - coords[:, 1].long() * x_stride + - coords[:, 2].long() * y_stride + - coords[:, 3].long() * z_stride) + flat_keys = coords[:, 0].long() * b_stride + flat_keys.add_(coords[:, 1], alpha=x_stride) + flat_keys.add_(coords[:, 2], alpha=y_stride) + flat_keys.add_(coords[:, 3], alpha=z_stride) vals = torch.arange(coords.shape[0], dtype=torch.int32, device=device) hashmap = TorchHashMap(flat_keys, vals) - neighbor = build_submanifold_neighbor_map( - hashmap, coords, W, H, D, Kw, Kh, Kd, - dilation[0], dilation[1], dilation[2] - ) + if cache_neighbor_map: + neighbor = build_submanifold_neighbor_map( + hashmap, coords, W, H, D, Kw, Kh, Kd, + dilation[0], dilation[1], dilation[2] + ) + else: + neighbor = None else: neighbor = neighbor_cache @@ -129,9 +135,6 @@ def sparse_submanifold_conv3d( output = torch.empty(N_pts, Co, device=device, dtype=feats.dtype) - # Zero row at index N_pts; missing neighbors (-1) gather it -> no separate masking. - feats_padded = torch.cat([feats, feats.new_zeros(1, Ci)], dim=0) - # Chunk over voxels to bound the (chunk, V, Ci) gather. max_chunk_mem_gb = get_recommended_chunk_mem(device) mem_per_row = V * Ci * feats.element_size() @@ -143,8 +146,17 @@ def sparse_submanifold_conv3d( end = min(start + chunk_size, N_pts) actual_chunk = end - start - chunk_idx = torch.where(neighbor[start:end] < 0, N_pts, neighbor[start:end]) # -1 -> zero row - gathered = feats_padded[chunk_idx] # (chunk, V, Ci) + if neighbor is None: + neighbor_chunk = build_submanifold_neighbor_map( + hashmap, coords[start:end], W, H, D, Kw, Kh, Kd, + dilation[0], dilation[1], dilation[2] + ) + else: + neighbor_chunk = neighbor[start:end] + + chunk_idx = neighbor_chunk.clamp_min(0) + gathered = feats[chunk_idx] # (chunk, V, Ci) + gathered.masked_fill_(neighbor_chunk[:, :, None] < 0, 0) gathered_flat = gathered.view(actual_chunk, V * Ci) output[start:end] = torch.matmul(gathered_flat, weight_T) # (chunk, V*Ci) @ (V*Ci, Co) diff --git a/comfy/ldm/trellis2/vae.py b/comfy/ldm/trellis2/vae.py index d0d00ce04..5a3b97a5c 100644 --- a/comfy/ldm/trellis2/vae.py +++ b/comfy/ldm/trellis2/vae.py @@ -33,6 +33,7 @@ def sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1, d self.kernel_size = tuple(kernel_size) if isinstance(kernel_size, (list, tuple)) else (kernel_size, ) * 3 self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, ) * 3 self.dilation = tuple(dilation) if isinstance(dilation, (list, tuple)) else (dilation, ) * 3 + self.cache_neighbor_map = True self.weight = nn.Parameter(torch.empty((out_channels, in_channels, *self.kernel_size))) if bias: @@ -60,10 +61,11 @@ def sparse_conv3d_forward(self, x): weight, bias, neighbor_cache, - self.dilation + self.dilation, + self.cache_neighbor_map, ) - if neighbor_cache is None: + if neighbor_cache is None and neighbor_cache_ is not None: x.register_spatial_cache(neighbor_cache_key, neighbor_cache_) out = x.replace(out) @@ -142,34 +144,39 @@ class SparseChannel2Spatial(nn.Module): super(SparseChannel2Spatial, self).__init__() self.factor = factor - def forward(self, x, subdivision = None): + def _subdivision_plan(self, x, subdivision): DIM = x.coords.shape[-1] - 1 - cache = x.get_spatial_cache(f'channel2spatial_{self.factor}') if cache is None: if subdivision is None: raise ValueError('Cache not found. Provide subdivision tensor or pair SparseChannel2Spatial with SparseSpatial2Channel.') - else: - sub = subdivision.feats # [N, self.factor ** DIM] - N_leaf = sub.sum(dim=-1) # [N] - subidx = sub.nonzero()[:, -1] - new_coords = x.coords.clone().detach() - new_coords[:, 1:] *= self.factor - new_coords = torch.repeat_interleave(new_coords, N_leaf, dim=0, output_size=subidx.shape[0]) - for i in range(DIM): - new_coords[:, i+1] += subidx // self.factor ** i % self.factor - idx = torch.repeat_interleave(torch.arange(x.coords.shape[0], device=x.device), N_leaf, dim=0, output_size=subidx.shape[0]) + flat_idx = subdivision.feats.flatten().nonzero(as_tuple=True)[0] + idx = torch.div(flat_idx, self.factor ** DIM, rounding_mode='floor') + subidx = flat_idx.remainder(self.factor ** DIM) + new_coords = x.coords[idx] + new_coords[:, 1:] *= self.factor + for i in range(DIM): + new_coords[:, i+1] += subidx // self.factor ** i % self.factor + return new_coords, flat_idx, False else: new_coords, idx, subidx = cache + flat_idx = idx * self.factor ** DIM + subidx + return new_coords, flat_idx, True + def _apply_plan(self, x, new_coords, flat_idx, inherit_cache): + DIM = x.coords.shape[-1] - 1 x_feats = x.feats.reshape(x.feats.shape[0] * self.factor ** DIM, -1) - new_feats = x_feats[idx * self.factor ** DIM + subidx] + new_feats = x_feats[flat_idx] out = SparseTensor(new_feats, new_coords, None if x._shape is None else torch.Size([x._shape[0], x._shape[1] // self.factor ** DIM])) out._scale = tuple([s / self.factor for s in x._scale]) - if cache is not None: # only keep cache when subdiv following it + if inherit_cache: # only keep cache when subdiv following it out._spatial_cache = dict(x._spatial_cache) return out + def forward(self, x, subdivision = None): + new_coords, flat_idx, inherit_cache = self._subdivision_plan(x, subdivision) + return self._apply_plan(x, new_coords, flat_idx, inherit_cache) + class SparseResBlockC2S3d(nn.Module): def __init__(self, channels: int, out_channels: Optional[int] = None, pred_subdiv: bool = True): super().__init__() @@ -192,13 +199,15 @@ class SparseResBlockC2S3d(nn.Module): h = h.replace(F.silu(h.feats, inplace=True)) h = self.conv1(h) subdiv_binarized = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None - h = self.updown(h, subdiv_binarized) - x = self.updown(x, subdiv_binarized) + new_coords, flat_idx, inherit_cache = self.updown._subdivision_plan(h, subdiv_binarized) + h = self.updown._apply_plan(h, new_coords, flat_idx, inherit_cache) + skip_feats = x.feats.reshape(x.feats.shape[0] * 8, -1)[flat_idx] + del x, new_coords, flat_idx, subdiv_binarized h = h.replace(self.norm2(h.feats)) h = h.replace(F.silu(h.feats, inplace=True)) h = self.conv2(h) skip_repeat = self.out_channels // (self.channels // 8) - h.feats.view(h.feats.shape[0], x.feats.shape[1], skip_repeat).add_(x.feats.to(h.feats.dtype).unsqueeze(-1)) + h.feats.view(h.feats.shape[0], skip_feats.shape[1], skip_repeat).add_(skip_feats.to(h.feats.dtype).unsqueeze(-1)) if self.pred_subdiv: return h, subdiv else: @@ -700,6 +709,9 @@ class SparseTensor(VarLenTensor): def __repr__(self) -> str: return f"SparseTensor(shape={self.shape}, dtype={self.dtype}, device={self.device})" + def _comfy_cache_tensors(self): + return self.data, self._spatial_cache + def sparse_cat(inputs: List[SparseTensor], dim: int = 0) -> SparseTensor: if dim == 0: start = 0 @@ -772,14 +784,15 @@ class SparseUnetVaeDecoder(nn.Module): ) ) if i < len(num_blocks) - 1: - self.blocks[-1].append( - globals()[up_block_type[i]]( - model_channels[i], - model_channels[i+1], - pred_subdiv=pred_subdiv, - **block_args[i], - ) + up_block = globals()[up_block_type[i]]( + model_channels[i], + model_channels[i+1], + pred_subdiv=pred_subdiv, + **block_args[i], ) + if num_blocks[i + 1] == 0: + up_block.conv2.cache_neighbor_map = False + self.blocks[-1].append(up_block) def forward(self, x: SparseTensor, guide_subs: Optional[List[SparseTensor]] = None, return_subs: bool = False) -> SparseTensor: h = self.from_latent(x) @@ -789,12 +802,17 @@ class SparseUnetVaeDecoder(nn.Module): if i < len(self.blocks) - 1 and j == len(res) - 1: if self.pred_subdiv: h, sub = block(h) - subs.append(sub) + subs.append(SparseTensor(feats=sub.feats, coords=sub.coords, shape=sub.shape, scale=sub._scale)) else: h = block(h, subdiv=guide_subs[i] if guide_subs is not None else None) else: h = block(h) - h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:])) + h = SparseTensor( + feats=F.layer_norm(h.feats, h.feats.shape[-1:]), + coords=h.coords, + shape=h.shape, + scale=h._scale, + ) h = self.output_layer(h) if return_subs: return h, subs diff --git a/comfy_api/latest/_util/geometry_types.py b/comfy_api/latest/_util/geometry_types.py index 64f49a97b..f94977117 100644 --- a/comfy_api/latest/_util/geometry_types.py +++ b/comfy_api/latest/_util/geometry_types.py @@ -12,6 +12,9 @@ class VOXEL: self.voxel_colors = voxel_colors self.resolution = resolution # each 3d model has its own resolution + def _comfy_cache_tensors(self): + return self.data, self.voxel_colors + class SPLAT: """A batch of 3D Gaussian splats in render-ready (activated, world-space) form. @@ -29,6 +32,9 @@ class SPLAT: self.sh = sh # (B, N, K, 3) spherical-harmonic color coefficients self.counts = counts # (B,) real lengths, or None + def _comfy_cache_tensors(self): + return self.positions, self.scales, self.rotations, self.opacities, self.sh, self.counts + class MESH: def __init__(self, vertices: torch.Tensor, faces: torch.Tensor, @@ -72,6 +78,23 @@ class MESH: self.material = material # SetMeshMaterial scalar/factor overrides self.emissive = emissive # emissive map: (B, H, W, 3) + def _comfy_cache_tensors(self): + return ( + self.vertices, + self.faces, + self.uvs, + self.vertex_colors, + self.texture, + self.metallic_roughness, + self.vertex_counts, + self.face_counts, + self.normals, + self.tangents, + self.normal_map, + self.material, + self.emissive, + ) + class File3D: """Class representing a 3D file from a file path or binary stream. diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index 6003bc7bd..9ff52fc48 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -564,19 +564,30 @@ class RAMPressureCache(LRUCache): ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE oom_ram_usage = ram_usage + seen_storages = set() def scan_list_for_ram_usage(outputs): nonlocal ram_usage, oom_ram_usage if outputs is None: return + if isinstance(outputs, Mapping): + outputs = outputs.values() + elif not isinstance(outputs, (list, tuple)): + outputs = (outputs,) for output in outputs: - if isinstance(output, (list, tuple)): + if isinstance(output, (list, tuple, Mapping)): scan_list_for_ram_usage(output) elif isinstance(output, torch.Tensor) and output.device.type == 'cpu': - ram_usage += output.numel() * output.element_size() - oom_ram_usage += output.numel() * output.element_size() + storage = output.untyped_storage() + storage_key = (storage.data_ptr(), storage.nbytes()) + if storage_key not in seen_storages: + seen_storages.add(storage_key) + ram_usage += storage.nbytes() + oom_ram_usage += storage.nbytes() elif is_model_patcher_output(output) and self.used_generation[key] != self.generation: #old ModelPatchers are the first to go oom_ram_usage = 1e30 + elif hasattr(output, "_comfy_cache_tensors"): + scan_list_for_ram_usage(output._comfy_cache_tensors()) scan_list_for_ram_usage(cache_entry.outputs) if ram_usage < min_entry_size: diff --git a/comfy_extras/mesh3d/uv_unwrap/pack.py b/comfy_extras/mesh3d/uv_unwrap/pack.py index 44103f8b6..6f238c5cd 100644 --- a/comfy_extras/mesh3d/uv_unwrap/pack.py +++ b/comfy_extras/mesh3d/uv_unwrap/pack.py @@ -27,6 +27,8 @@ except ImportError: # Cap on deterministic sweep density: tiny charts on a large atlas would otherwise enumerate every texel column. _SWEEP_CAP = 1024 +_TORCH_PREP_MAX_ANGLE_ELEMENTS = 1 << 23 +_TORCH_RASTER_FACE_BATCH = 1 << 17 @dataclass @@ -429,78 +431,83 @@ def _dilate_local(x: Tensor, p: int) -> Tensor: return x -def _raster_all_torch(uvs_tex_pad, faces_pad, fmask, bw_t, bh_t, padding, device): +def _raster_all_torch(uvs_tex, uv_offsets, faces_cat, face_offsets, bw_t, bh_t, padding, device): """Rasterize every chart into one flat bool buffer; buf[cbase[i]:cbase[i+1]].view(bh,bw) is chart i's bitmap. Triangles are bucketed by next-pow2 bbox size to bound memory.""" - n = uvs_tex_pad.shape[0] - fmax = faces_pad.shape[1] + n = len(face_offsets) - 1 bwL, bhL = bw_t.long(), bh_t.long() cbase = torch.zeros(n + 1, dtype=torch.long, device=device) torch.cumsum(bwL * bhL, 0, out=cbase[1:]) buf = torch.zeros(int(cbase[-1].item()), dtype=torch.bool, device=device) - # gather all triangle coords, keep only valid faces -> (Ttot,3,2) + chart id per triangle - fp = faces_pad.reshape(n, fmax * 3) - tri = torch.gather(uvs_tex_pad, 1, fp[..., None].expand(-1, -1, 2)).reshape(n * fmax, 3, 2) - fm = fmask.reshape(-1) - tri_f = tri[fm] - if tri_f.shape[0] == 0: + total_faces = int(faces_cat.shape[0]) + if total_faces == 0: return buf, cbase - cid = torch.arange(n, device=device).repeat_interleave(fmax)[fm] - - # per-triangle pixel bbox, inflated by padding (origin >= 0); bucket by next-pow2 max-dim - tmin = tri_f.amin(1) - tmax = tri_f.amax(1) - x0 = (tmin[:, 0].floor().long() - padding).clamp_min(0) - y0 = (tmin[:, 1].floor().long() - padding).clamp_min(0) - bbw = (tmax[:, 0].ceil().long() + padding) - x0 + 1 - bbh = (tmax[:, 1].ceil().long() + padding) - y0 + 1 - mxd = torch.maximum(bbw, bbh).clamp_min(1) - bsz = (2 ** torch.ceil(torch.log2(mxd.float())).long()).long() - - a = tri_f[:, 0] - b = tri_f[:, 1] - c = tri_f[:, 2] - v0 = b - a - v1 = c - a - d00 = (v0 * v0).sum(-1) - d01 = (v0 * v1).sum(-1) - d11 = (v1 * v1).sum(-1) - den = (d00 * d11 - d01 * d01).clamp(min=1e-20) - - + faces = torch.from_numpy(np.ascontiguousarray(faces_cat)).to(device=device, dtype=torch.long) + face_counts = torch.from_numpy(np.diff(face_offsets)).to(device=device, dtype=torch.long) + face_charts = torch.arange(n, device=device).repeat_interleave(face_counts) + uv_offsets_t = torch.as_tensor(uv_offsets, dtype=torch.long, device=device) free = comfy.model_management.get_free_memory(device) budget = int(min(1 << 23, max(1 << 20, (free * 0.25) / 56))) - for g in sorted(set(bsz.tolist())): # one batch per pow2 grid - sel_g = (bsz == g).nonzero(as_tuple=True)[0] - per = max(1, budget // (g * g)) - for cs in range(0, sel_g.shape[0], per): - sel = sel_g[cs:cs + per] - m = sel.shape[0] - xs0 = x0[sel].view(m, 1, 1) - ys0 = y0[sel].view(m, 1, 1) - cc = cid[sel] - bwp = bwL[cc].view(m, 1, 1) - bhp = bhL[cc].view(m, 1, 1) - gi = torch.arange(g, device=device) - px = xs0 + gi.view(1, 1, g) - py = ys0 + gi.view(1, g, 1) # (m,g,g) int - pxf = px.float() + 0.5 - pyf = py.float() + 0.5 - v2x = pxf - a[sel, 0].view(m, 1, 1) - v2y = pyf - a[sel, 1].view(m, 1, 1) - d20 = v2x * v0[sel, 0].view(m, 1, 1) + v2y * v0[sel, 1].view(m, 1, 1) - d21 = v2x * v1[sel, 0].view(m, 1, 1) + v2y * v1[sel, 1].view(m, 1, 1) - idn = den[sel].view(m, 1, 1).reciprocal() - vv = torch.addcmul(d11[sel].view(m, 1, 1) * d20, d01[sel].view(m, 1, 1), d21, value=-1) * idn - ww = torch.addcmul(d00[sel].view(m, 1, 1) * d21, d01[sel].view(m, 1, 1), d20, value=-1) * idn - uu = 1.0 - vv - ww - inside = (uu >= -1e-6) & (vv >= -1e-6) & (ww >= -1e-6) - if padding > 0: - inside = _dilate_local(inside, padding) - valid = inside & (px < bwp) & (py < bhp) - flat = (cbase[cc].view(m, 1, 1) + py * bwp + px)[valid] - buf[flat] = True + for fs in range(0, total_faces, _TORCH_RASTER_FACE_BATCH): + fe = min(fs + _TORCH_RASTER_FACE_BATCH, total_faces) + cid = face_charts[fs:fe] + tri_f = uvs_tex[faces[fs:fe] + uv_offsets_t[cid, None]] + + # per-triangle pixel bbox, inflated by padding (origin >= 0); bucket by next-pow2 max-dim + tmin = tri_f.amin(1) + tmax = tri_f.amax(1) + x0 = (tmin[:, 0].floor().long() - padding).clamp_min(0) + y0 = (tmin[:, 1].floor().long() - padding).clamp_min(0) + bbw = (tmax[:, 0].ceil().long() + padding) - x0 + 1 + bbh = (tmax[:, 1].ceil().long() + padding) - y0 + 1 + mxd = torch.maximum(bbw, bbh).clamp_min(1) + bsz = (2 ** torch.ceil(torch.log2(mxd.float())).long()).long() + + a = tri_f[:, 0] + b = tri_f[:, 1] + c = tri_f[:, 2] + v0 = b - a + v1 = c - a + d00 = (v0 * v0).sum(-1) + d01 = (v0 * v1).sum(-1) + d11 = (v1 * v1).sum(-1) + den = (d00 * d11 - d01 * d01).clamp(min=1e-20) + + for g in torch.unique(bsz).tolist(): # one batch per pow2 grid + sel_g = (bsz == g).nonzero(as_tuple=True)[0] + per = max(1, budget // (g * g)) + for cs in range(0, sel_g.shape[0], per): + sel = sel_g[cs:cs + per] + m = sel.shape[0] + xs0 = x0[sel].view(m, 1, 1) + ys0 = y0[sel].view(m, 1, 1) + cc = cid[sel] + bwp = bwL[cc].view(m, 1, 1) + bhp = bhL[cc].view(m, 1, 1) + gi = torch.arange(g, device=device) + px = xs0 + gi.view(1, 1, g) + py = ys0 + gi.view(1, g, 1) # (m,g,g) int + pxf = px.float() + 0.5 + pyf = py.float() + 0.5 + v2x = pxf - a[sel, 0].view(m, 1, 1) + v2y = pyf - a[sel, 1].view(m, 1, 1) + d20 = v2x * v0[sel, 0].view(m, 1, 1) + v2y * v0[sel, 1].view(m, 1, 1) + d21 = v2x * v1[sel, 0].view(m, 1, 1) + v2y * v1[sel, 1].view(m, 1, 1) + idn = den[sel].view(m, 1, 1).reciprocal() + vv = torch.addcmul(d11[sel].view(m, 1, 1) * d20, d01[sel].view(m, 1, 1), d21, value=-1) * idn + ww = torch.addcmul(d00[sel].view(m, 1, 1) * d21, d01[sel].view(m, 1, 1), d20, value=-1) * idn + uu = 1.0 - vv - ww + inside = (uu >= -1e-6) & (vv >= -1e-6) & (ww >= -1e-6) + if padding > 0: + inside = _dilate_local(inside, padding) + valid = inside & (px < bwp) & (py < bhp) + flat = (cbase[cc].view(m, 1, 1) + py * bwp + px)[valid] + buf[flat] = True + del sel, xs0, ys0, cc, bwp, bhp, gi, px, py, pxf, pyf + del v2x, v2y, d20, d21, idn, vv, ww, uu, inside, valid, flat + del cid, tri_f, tmin, tmax, x0, y0, bbw, bbh, mxd, bsz + del a, b, c, v0, v1, d00, d01, d11, den return buf, cbase @@ -594,65 +601,99 @@ def _best_placement_torch(atlas, pix0, dim0, dim1, cands, n_sky, cur_w, cur_h, d return sky -def _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces, +def _pack_bitmap_torch(uvs_cat, uv_offsets, chart_3d_areas, chart_uv_areas, faces_cat, face_offsets, texels_per_unit, padding_texels, attempts=4096, rng_seed=0, progress_callback=None): """Torch rasterize-and-place packer (numba-free fallback). Returns (placements, atlas_w, atlas_h).""" - n = len(chart_uvs) + n = len(uv_offsets) - 1 if n == 0: return [], 1, 1 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() - # ---- Prepare pass 1: best-rotation + scale + bbox for ALL charts at once (batched) ---- - vcount = [int(u.shape[0]) for u in chart_uvs] - fcount = [int(f.shape[0]) for f in chart_faces] - vmax = max(vcount) - fmax = max(fcount) - uvs_pad = torch.zeros(n, vmax, 2, device=device) - vmask = torch.zeros(n, vmax, dtype=torch.bool, device=device) - faces_pad = torch.zeros(n, fmax, 3, dtype=torch.long, device=device) - fmask = torch.zeros(n, fmax, dtype=torch.bool, device=device) - for i in range(n): - uvs_pad[i, :vcount[i]] = chart_uvs[i].to(device=device, dtype=torch.float32) - vmask[i, :vcount[i]] = True - if fcount[i]: - faces_pad[i, :fcount[i]] = chart_faces[i].to(device=device, dtype=torch.long) - fmask[i, :fcount[i]] = True - u0, u1 = uvs_pad[..., 0], uvs_pad[..., 1] # (N,Vmax) - BIG = 1e30 - mlo = torch.where(vmask, torch.zeros_like(u0), u0.new_full((), BIG)) - mhi = torch.where(vmask, torch.zeros_like(u0), u0.new_full((), -BIG)) - xr = torch.addcmul(u0[:, :, None] * cos_a, u1[:, :, None], sin_a, value=-1) # (N,Vmax,A) - yr = torch.addcmul(u0[:, :, None] * sin_a, u1[:, :, None], cos_a) - xsp = (xr + mhi[:, :, None]).amax(1) - (xr + mlo[:, :, None]).amin(1) # (N,A) masked span - ysp = (yr + mhi[:, :, None]).amax(1) - (yr + mlo[:, :, None]).amin(1) - ti = (xsp * ysp).argmin(1) # (N,) best angle per chart - cc, ss = cos_a[ti][:, None], sin_a[ti][:, None] # (N,1) - rx = torch.addcmul(u0 * cc, u1, ss, value=-1) # (N,Vmax) - ry = torch.addcmul(u0 * ss, u1, cc) - rxmin = (rx + mlo).amin(1) # (N,) - rxmax = (rx + mhi).amax(1) - rymin = (ry + mlo).amin(1) - rymax = (ry + mhi).amax(1) + # ---- Prepare pass 1: best-rotation + scale + bbox in bounded batches ---- + vcount = np.diff(uv_offsets).astype(np.int64, copy=False) + uvs_src = torch.from_numpy(np.ascontiguousarray(uvs_cat)).to(device=device, dtype=torch.float32) + uvs_tex = torch.empty_like(uvs_src) a3 = torch.tensor([max(a, 1e-12) for a in chart_3d_areas], device=device) au = torch.tensor([max(a, 1e-12) for a in chart_uv_areas], device=device) - base = (a3 / au).sqrt() * texels_per_unit - maxb = (4.0 * a3.sqrt() * texels_per_unit).clamp_min(8.0) - bbm = torch.maximum(rxmax - rxmin, rymax - rymin).clamp_min(1e-12) - scale = torch.minimum(base, maxb / bbm) # (N,) - uvs_tex_pad = torch.stack([(rx - rxmin[:, None]) * scale[:, None], - (ry - rymin[:, None]) * scale[:, None]], dim=-1) # (N,Vmax,2) - bw_t = ((rxmax - rxmin) * scale).ceil().int() + padding_texels + 1 - bh_t = ((rymax - rymin) * scale).ceil().int() + padding_texels + 1 + ti_all = torch.empty(n, dtype=torch.long, device=device) + scale_all = torch.empty(n, device=device) + bw_t = torch.empty(n, dtype=torch.int, device=device) + bh_t = torch.empty(n, dtype=torch.int, device=device) + start = 0 + while start < n: + end = start + vmax = 0 + while end < n: + next_vmax = max(vmax, int(vcount[end])) + if end > start and (end - start + 1) * next_vmax * 36 > _TORCH_PREP_MAX_ANGLE_ELEMENTS: + break + vmax = next_vmax + end += 1 + m = end - start + uvs_pad = torch.zeros(m, vmax, 2, device=device) + vmask = torch.zeros(m, vmax, dtype=torch.bool, device=device) + for j, i in enumerate(range(start, end)): + count = int(vcount[i]) + uvs_pad[j, :count] = uvs_src[int(uv_offsets[i]):int(uv_offsets[i + 1])] + vmask[j, :count] = True + u0, u1 = uvs_pad[..., 0], uvs_pad[..., 1] + xsp = torch.empty(m, 36, device=device) + ysp = torch.empty(m, 36, device=device) + angle_batch = max(1, min(36, _TORCH_PREP_MAX_ANGLE_ELEMENTS // max(m * vmax, 1))) + mask3 = vmask[:, :, None] + for a0 in range(0, 36, angle_batch): + a1 = min(a0 + angle_batch, 36) + cc = cos_a[a0:a1] + ss = sin_a[a0:a1] + xr = torch.addcmul(u0[:, :, None] * cc, u1[:, :, None], ss, value=-1) + xr.masked_fill_(~mask3, -1e30) + xmax = xr.amax(1) + xr.masked_fill_(~mask3, 1e30) + xsp[:, a0:a1] = xmax - xr.amin(1) + del xr + yr = torch.addcmul(u0[:, :, None] * ss, u1[:, :, None], cc) + yr.masked_fill_(~mask3, -1e30) + ymax = yr.amax(1) + yr.masked_fill_(~mask3, 1e30) + ysp[:, a0:a1] = ymax - yr.amin(1) + del yr + ti = (xsp * ysp).argmin(1) + cc, ss = cos_a[ti][:, None], sin_a[ti][:, None] + rx = torch.addcmul(u0 * cc, u1, ss, value=-1) + ry = torch.addcmul(u0 * ss, u1, cc) + rx.masked_fill_(~vmask, -1e30) + rxmax = rx.amax(1) + rx.masked_fill_(~vmask, 1e30) + rxmin = rx.amin(1) + ry.masked_fill_(~vmask, -1e30) + rymax = ry.amax(1) + ry.masked_fill_(~vmask, 1e30) + rymin = ry.amin(1) + base = (a3[start:end] / au[start:end]).sqrt() * texels_per_unit + maxb = (4.0 * a3[start:end].sqrt() * texels_per_unit).clamp_min(8.0) + bbm = torch.maximum(rxmax - rxmin, rymax - rymin).clamp_min(1e-12) + scale = torch.minimum(base, maxb / bbm) + uvs_batch = torch.stack([(rx - rxmin[:, None]) * scale[:, None], + (ry - rymin[:, None]) * scale[:, None]], dim=-1) + uvs_tex[int(uv_offsets[start]):int(uv_offsets[end])] = uvs_batch[vmask] + ti_all[start:end] = ti + scale_all[start:end] = scale + bw_t[start:end] = ((rxmax - rxmin) * scale).ceil().int() + padding_texels + 1 + bh_t[start:end] = ((rymax - rymin) * scale).ceil().int() + padding_texels + 1 + del uvs_pad, vmask, u0, u1, xsp, ysp, mask3, rx, ry, uvs_batch + start = end # one sync: pull all per-chart scalars - thetas = ang[ti].cpu().tolist() - scales = scale.cpu().tolist() + thetas = ang[ti_all].cpu().tolist() + scales = scale_all.cpu().tolist() + del uvs_src, a3, au, ti_all, scale_all # ---- Prepare pass 2: rasterize ALL charts at once, then derive per-chart sparse data ---- - buf, cbase = _raster_all_torch(uvs_tex_pad, faces_pad, fmask, bw_t, bh_t, padding_texels, device) + buf, cbase = _raster_all_torch( + uvs_tex, uv_offsets, faces_cat, face_offsets, bw_t, bh_t, padding_texels, device) # nonzero over the flat buffer is ascending, so pixels come out grouped by chart nz = buf.nonzero(as_tuple=True)[0] @@ -679,14 +720,6 @@ def _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces, pix_l = [pix_all[offs[i]:offs[i + 1]] for i in range(n)] pixr_l = [pixr_all[offs[i]:offs[i + 1]] for i in range(n)] - # column tops (skyline lift), batched via flat scatter-amax over (chart, column) keys - wmax = max(max(h, w) for (h, w) in dim_l) - ct_pad = torch.full((n * wmax,), -1, dtype=torch.long, device=device) - ctr_pad = torch.full((n * wmax,), -1, dtype=torch.long, device=device) - ct_pad.scatter_reduce_(0, cid * wmax + px, py, reduce="amax") - ctr_pad.scatter_reduce_(0, cid * wmax + (rmax[cid] - py), px, reduce="amax") - ct_pad = ct_pad.view(n, wmax) - ctr_pad = ctr_pad.view(n, wmax) del cid, py, px, rmax, cmax # ---- Placement: skyline bin-pack on GPU ---- @@ -731,7 +764,8 @@ def _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces, atlas[by + pix[:, 0], bx + pix[:, 1]] = True # sparse blit cur_w = max(cur_w, bx + bw_) cur_h = max(cur_h, by + bh_) - ct = (ctr_pad if swap else ct_pad)[ci, :bw_] # GPU skyline lift + ct = torch.full((bw_,), -1, dtype=torch.long, device=device) + ct.scatter_reduce_(0, pix[:, 1], pix[:, 0], reduce="amax") ix = ar[bx:bx + bw_] sky_t[ix] = torch.where(ct >= 0, torch.maximum(sky_t[ix], by + ct + 1), sky_t[ix]) placements[ci] = ChartPlacement(chart_id=ci, offset=(float(bx), float(by)), @@ -761,13 +795,9 @@ def pack_bitmap_concat( if n == 0: return empty, empty, empty, empty.astype(np.float64), empty.astype(np.float64), empty, 1, 1 if not _HAVE_NUMBA_PACK: - chart_uvs = [torch.from_numpy(np.ascontiguousarray(uvs_cat[uv_offsets[c]:uv_offsets[c + 1]])) - for c in range(n)] - chart_faces = [torch.from_numpy(np.ascontiguousarray(faces_cat[face_offsets[c]:face_offsets[c + 1]])) - for c in range(n)] placements, w, h = _pack_bitmap_torch( - chart_uvs, [float(a) for a in chart_3d_areas], [float(a) for a in chart_uv_areas], - chart_faces, texels_per_unit, padding_texels, attempts=attempts, + uvs_cat, uv_offsets, [float(a) for a in chart_3d_areas], [float(a) for a in chart_uv_areas], + faces_cat, face_offsets, texels_per_unit, padding_texels, attempts=attempts, rng_seed=rng_seed, progress_callback=progress_callback) px = np.array([p.offset[0] for p in placements], dtype=np.int64) py = np.array([p.offset[1] for p in placements], dtype=np.int64) diff --git a/comfy_extras/nodes_mesh_postprocess.py b/comfy_extras/nodes_mesh_postprocess.py index 99725efd5..d730cbff9 100644 --- a/comfy_extras/nodes_mesh_postprocess.py +++ b/comfy_extras/nodes_mesh_postprocess.py @@ -23,6 +23,21 @@ from scipy.spatial import cKDTree import scipy.ndimage as ndi +def _mesh_face_count(mesh): + if mesh.face_counts is not None: + return sum(int(count) for count in mesh.face_counts) + if isinstance(mesh.faces, list): + return sum(int(faces.shape[0]) for faces in mesh.faces) + return int(mesh.faces.numel() // 3) + + +def _prepare_gpu_mesh_processing(device, memory_required): + comfy.model_management.free_memory( + int(memory_required) + comfy.model_management.minimum_inference_memory(), + device, + ) + + def paint_mesh_with_voxels(mesh, voxel_coords, voxel_colors, resolution): """Paint a mesh using nearest-neighbor colors from a sparse voxel field.""" device = comfy.model_management.vae_offload_device() @@ -215,6 +230,7 @@ def _bake_position_map(verts_np, faces_np, uvs_np, texture_size): verts = torch.from_numpy(np.ascontiguousarray(verts_np, dtype=np.float32)).to(dev) faces = torch.from_numpy(np.ascontiguousarray(faces_np).astype(np.int64)).to(dev) attr = _interp_vertex_attr(verts, faces, face_idx, bary, mask) + del face_idx, bary, verts, faces return attr.cpu().numpy(), mask.cpu().numpy() @@ -270,14 +286,15 @@ def _trilinear_sample_sparse_gpu(positions, voxel_coords_np, color_np, resolutio R = int(resolution) origin = -0.5 voxel_size = 1.0 / R + index_dtype = torch.int32 if R ** 3 <= torch.iinfo(torch.int32).max else torch.int64 P = torch.from_numpy(np.ascontiguousarray(positions)).to(dev).float() - VC = torch.from_numpy(np.ascontiguousarray(voxel_coords_np)).to(dev).long() + VC = torch.from_numpy(np.ascontiguousarray(voxel_coords_np)).to(device=dev, dtype=index_dtype) col = torch.from_numpy(np.ascontiguousarray(color_np)).to(dev).float() K, C = P.shape[0], col.shape[1] M = VC.shape[0] # Cell-CENTER convention (see NumPy path): -0.5 to bracket the query. gc = (P - origin) / voxel_size - 0.5 - base = torch.floor(gc).long() + base = torch.floor(gc).to(index_dtype) frac = gc - base.float() key = (VC[:, 0] * R + VC[:, 1]) * R + VC[:, 2] skey, order = key.sort() @@ -298,7 +315,9 @@ def _trilinear_sample_sparse_gpu(positions, voxel_coords_np, color_np, resolutio matched = inb & (skey[ins] == qk) idx = order[ins] # garbage where !matched w = torch.where(matched, wx * wy * wz, torch.zeros_like(wx))[:, None] - acc += w * col[idx] # w=0 cancels garbage rows + weighted = col[idx] + weighted.mul_(w) # w=0 cancels garbage rows + acc.add_(weighted) wsum += w ok = wsum[:, 0] > 1e-8 vals = torch.zeros((K, C), device=dev) @@ -316,8 +335,9 @@ def _nearest_voxel_sample_gpu(positions, voxel_coords_np, color_np, resolution): float32, found [K] bool); `found` is False for stragglers left to the caller's cKDTree.""" dev = comfy.model_management.get_torch_device() R = int(resolution) + index_dtype = torch.int32 if R ** 3 <= torch.iinfo(torch.int32).max else torch.int64 P = torch.from_numpy(np.ascontiguousarray(positions)).to(dev).float() - VC = torch.from_numpy(np.ascontiguousarray(voxel_coords_np)).to(dev).long() + VC = torch.from_numpy(np.ascontiguousarray(voxel_coords_np)).to(device=dev, dtype=index_dtype) col = torch.from_numpy(np.ascontiguousarray(color_np)).to(dev).float() M, K = VC.shape[0], P.shape[0] key = (VC[:, 0] * R + VC[:, 1]) * R + VC[:, 2] @@ -327,7 +347,7 @@ def _nearest_voxel_sample_gpu(positions, voxel_coords_np, color_np, resolution): """Nearest occupied voxel within ±radius cells, for query subset P[idx].""" Ps = P[idx] # Cell-CENTER convention: nearest coord = round((p+0.5)*R-0.5) (matches official). - rc = ((Ps + 0.5) * R - 0.5).round().long() + rc = ((Ps + 0.5) * R - 0.5).round().to(index_dtype) n = idx.shape[0] bd = torch.full((n,), 1e30, device=dev) bi = torch.zeros(n, dtype=torch.long, device=dev) @@ -336,7 +356,7 @@ def _nearest_voxel_sample_gpu(positions, voxel_coords_np, color_np, resolution): for dx in rng: for dy in rng: for dz in rng: - cc = rc + torch.tensor([dx, dy, dz], device=dev) + cc = rc + torch.tensor([dx, dy, dz], dtype=index_dtype, device=dev) inb = ((cc >= 0) & (cc < R)).all(1) qk = (cc[:, 0] * R + cc[:, 1]) * R + cc[:, 2] ins = torch.searchsorted(skey, qk).clamp(max=M - 1) @@ -401,12 +421,7 @@ def _sample_voxel_attrs_per_texel(position_map, mask, voxel_coords, voxel_colors if not mask.any(): return out - 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 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): @@ -415,6 +430,9 @@ def _sample_voxel_attrs_per_texel(position_map, mask, voxel_coords, voxel_colors # 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(): + origin = np.array([-0.5, -0.5, -0.5], dtype=np.float32) + voxel_size = 1.0 / float(resolution) + voxel_pos = (coords_np.astype(np.float32) + 0.5) * voxel_size + origin tree = cKDTree(voxel_pos) _, nearest_idx = tree.query(query[~found], k=1, workers=-1) vals[~found] = color_np[nearest_idx] @@ -428,7 +446,8 @@ def _sample_voxel_attrs_per_texel(position_map, mask, voxel_coords, voxel_colors vals, ok = _trilinear_sample_sparse(valid_positions, coords_np, color_np, resolution) if not ok.all(): vals[~ok] = _nearest(valid_positions[~ok]) # no occupied neighbour - out[mask] = np.clip(vals, 0.0, 1.0).astype(np.float32) + np.clip(vals, 0.0, 1.0, out=vals) + out[mask] = vals return out @@ -470,8 +489,9 @@ def _build_triangle_bvh(tri): span = (hi - lo).clamp_min(1e-12) q = (((cent - lo) / span).clamp(0, 1) * float((1 << 21) - 1)).long() morton = (_morton_expand21(q[:, 0]) << 2 | _morton_expand21(q[:, 1]) << 1 | _morton_expand21(q[:, 2])).long() - order = torch.argsort(morton) - msort = morton[order] + order_long = torch.argsort(morton) + msort = morton[order_long] + order = order_long.to(torch.int32) # delta(i,j): common-prefix length of (morton, index) keys of leaves i,j (index # breaks ties so duplicate codes still split); -1 if OOB. @@ -485,7 +505,7 @@ def _build_triangle_bvh(tri): cpi = torch.where(xi == 0, torch.full_like(x, 32), 31 - _msb_int64(xi.clamp_min(1))) return torch.where(ok, cp + torch.where(same, cpi, torch.zeros_like(cp)), torch.full_like(x, -1)) - I = torch.arange(T - 1, device=dev) + I = torch.arange(T - 1, dtype=torch.int32, device=dev) dplus = delta(I, I + 1) dminus = delta(I, I - 1) direction = torch.where(dplus >= dminus, torch.ones_like(I), -torch.ones_like(I)) @@ -529,11 +549,18 @@ def _build_triangle_bvh(tri): left = torch.where(gamma == first, LEAF + gamma, gamma) right = torch.where(gamma + 1 == last, LEAF + gamma + 1, gamma + 1) + del cent, lo, hi, span, q, morton, order_long + del I, dplus, dminus, direction, dmin, lmax, cond, l, t, j + del first, last, dnode, s, div, rng, step, cond1, gamma + delta = None + msort = None + # node AABBs: leaves seeded, internal unioned bottom-up (~log2(T) passes; cap is a backstop). nmin = torch.empty((2 * T, 3), device=dev) nmax = torch.empty((2 * T, 3), device=dev) nmin[LEAF:] = amin[order] nmax[LEAF:] = amax[order] + del amin, amax setm = torch.zeros(2 * T, dtype=torch.bool, device=dev) setm[LEAF:] = True for _ in range(128): @@ -563,7 +590,7 @@ def _closest_points_on_mesh_bvh(Q, tri, bvh, max_stack=64, return_face=False): left = bvh['left'] right = bvh['right'] order = bvh['order'] - stack = torch.full((N, max_stack), -1, dtype=torch.long, device=dev) + stack = torch.full((N, max_stack), -1, dtype=torch.int32, device=dev) sp = torch.ones(N, dtype=torch.long, device=dev) stack[:, 0] = 0 best = torch.full((N,), 1e30, device=dev) @@ -593,7 +620,7 @@ def _closest_points_on_mesh_bvh(Q, tri, bvh, max_stack=64, return_face=False): gu = ga[upd] best[gu] = d2[upd] bestp[gu] = cp[upd] - bestf[gu] = fidx[upd] + bestf[gu] = fidx[upd].long() iv = within & ~isleaf if bool(iv.any()): gi = a[iv] @@ -631,20 +658,19 @@ def _back_project_positions(position_map, mask, ref_v, ref_f, max_query_res=1024 dev = comfy.model_management.get_torch_device() rv = ref_v.detach().to(dev).float() - rf = ref_f.detach().to(dev).long() + rf = ref_f.detach().to(device=dev, dtype=torch.int32) tri = rv[rf] bvh = _build_triangle_bvh(tri) - - def _closest(pts_np): - return _closest_points_on_mesh_bvh( - torch.from_numpy(np.ascontiguousarray(pts_np.astype(np.float32))).to(dev), tri, bvh - ).detach().cpu().numpy().astype(np.float32) + del rv, rf H, W, _ = position_map.shape stride = max(1, int(math.ceil(max(H, W) / float(max_query_res)))) if stride == 1 or not mask[::stride, ::stride].any(): out = position_map.copy() - out[mask] = _closest(position_map[mask]).astype(position_map.dtype) + query = torch.from_numpy(np.ascontiguousarray(position_map[mask], dtype=np.float32)).to(dev) + closest = _closest_points_on_mesh_bvh(query, tri, bvh).cpu().numpy().astype(position_map.dtype) + del query, tri, bvh + out[mask] = closest return out # Low-res correction, then bilinear upsample to full resolution. @@ -652,7 +678,10 @@ def _back_project_positions(position_map, mask, ref_v, ref_f, max_query_res=1024 mask_lo = mask[::stride, ::stride] Hl, Wl = mask_lo.shape corr_lo = np.zeros((Hl, Wl, 3), dtype=np.float32) - corr_lo[mask_lo] = _closest(pos_lo[mask_lo]) - pos_lo[mask_lo].astype(np.float32) + query = torch.from_numpy(np.ascontiguousarray(pos_lo[mask_lo], dtype=np.float32)).to(dev) + closest = _closest_points_on_mesh_bvh(query, tri, bvh).cpu().numpy().astype(np.float32) + del query, tri, bvh + corr_lo[mask_lo] = closest - pos_lo[mask_lo].astype(np.float32) inds = ndi.distance_transform_edt(~mask_lo, return_distances=False, return_indices=True) corr_lo = corr_lo[tuple(inds)] # extrapolate into gutter (nearest) corr = torch.nn.functional.interpolate( @@ -782,7 +811,7 @@ def _closest_hit_rays_bvh(orig, dirs, tri, bvh, tmin=0.0, tmax=1e30, max_stack=6 upd = h & (t < best_t[ga]) gu = ga[upd] best_t[gu] = t[upd] - best_f[gu] = fidx[upd] + best_f[gu] = fidx[upd].long() iv = within & ~isleaf if bool(iv.any()): gi = a[iv] @@ -837,7 +866,7 @@ def _bake_ambient_occlusion(high_v, high_f, low_v_np, low_f_np, low_uv_np, low_n Nl = torch.nn.functional.normalize((bsel[:, :, None] * low_n[vtri]).sum(1), dim=-1, eps=1e-6) hv = high_v.to(dev).float() - hf = high_f.to(dev).long() + hf = high_f.to(device=dev, dtype=torch.int32) tri = hv[hf] bvh = _build_triangle_bvh(tri) diag = float((hv.amax(0) - hv.amin(0)).norm().clamp_min(1e-6)) @@ -1155,7 +1184,7 @@ def _bake_normal_map(high_v, high_f, high_n, low_v_np, low_f_np, low_uv_np, low_ Wl = _interp(tangents[:, 3:4])[:, 0] hv = high_v.to(dev).float() - hf = high_f.to(dev).long() + hf = high_f.to(device=dev, dtype=torch.int32) tri = hv[hf] bvh = _build_triangle_bvh(tri) @@ -1209,10 +1238,13 @@ def _jfa_fill_gpu(img01, mask): it = torch.from_numpy(np.ascontiguousarray(img01)).to(dev).float() mm = torch.from_numpy(np.ascontiguousarray(mask)).to(dev) H, W = mm.shape - yy, xx = torch.meshgrid(torch.arange(H, device=dev), torch.arange(W, device=dev), indexing="ij") + yy, xx = torch.meshgrid( + torch.arange(H, dtype=torch.int32, device=dev), + torch.arange(W, dtype=torch.int32, device=dev), + indexing="ij", + ) by = torch.where(mm, yy, torch.full_like(yy, -1)) bx = torch.where(mm, xx, torch.full_like(xx, -1)) - INF = torch.full_like(yy, 1 << 30) step = 1 << ((max(H, W) - 1).bit_length() - 1) while step >= 1: for dy in (-step, 0, step): @@ -1224,13 +1256,14 @@ def _jfa_fill_gpu(img01, mask): cby = by[ny, nx] cbx = bx[ny, nx] valid = cby >= 0 - dc = torch.where(valid, (yy - cby) ** 2 + (xx - cbx) ** 2, INF) - db = torch.where(by >= 0, (yy - by) ** 2 + (xx - bx) ** 2, INF) + dc = torch.where(valid, (yy - cby) ** 2 + (xx - cbx) ** 2, 1 << 30) + db = torch.where(by >= 0, (yy - by) ** 2 + (xx - bx) ** 2, 1 << 30) take = valid & (dc < db) by = torch.where(take, cby, by) bx = torch.where(take, cbx, bx) step //= 2 - filled = it[by.clamp(0).long(), bx.clamp(0).long()] + flat_idx = by.clamp_min_(0).long().mul_(W).add_(bx.clamp_min_(0)) + filled = it.reshape(-1, it.shape[-1])[flat_idx.reshape(-1)].reshape(H, W, -1) return filled.cpu().numpy() @@ -1309,18 +1342,22 @@ def bake_texture_from_voxel_fn(vertices, faces, voxel_coords, voxel_colors, # PBR layout (upstream pbr_attr_layout): 0:3 base_color, 3 metallic, 4 roughness, 5 alpha. C = attrs.shape[-1] - base_color = attrs[..., 0:3] + base_color = np.ascontiguousarray(attrs[..., 0:3]) has_pbr = C >= 5 - metallic = attrs[..., 3:4] if C >= 4 else None - roughness = attrs[..., 4:5] if C >= 5 else None # alpha (idx 5) ignored — meshes kept opaque (upstream OPAQUE alpha_mode). - base_color = _seam_fill(np.ascontiguousarray(base_color), mask) mr_image = None if has_pbr: # glTF metallicRoughness: R unused, G=roughness, B=metallic. - mr = np.concatenate([np.zeros_like(roughness), roughness, metallic], axis=-1) - mr_image = _seam_fill(np.ascontiguousarray(mr), mask) + mr = np.empty((*attrs.shape[:-1], 3), dtype=attrs.dtype) + mr[..., 0] = 0.0 + mr[..., 1] = attrs[..., 4] + mr[..., 2] = attrs[..., 3] + del attrs + + base_color = _seam_fill(base_color, mask) + if has_pbr: + mr_image = _seam_fill(mr, mask) device = vertices.device out_v = torch.from_numpy(new_verts).to(device=device, dtype=torch.float32) @@ -1338,9 +1375,9 @@ def _mr_channel(packed_mr, ch, ref): """Pull one channel (G=roughness idx 1, B=metallic idx 2) from a packed glTF MR map as 3-channel grayscale [H,W,3] in [0,1]. Black sized like `ref` if no MR map.""" if packed_mr is None: - return torch.zeros_like(ref.float().cpu()) - m = packed_mr.float().clamp(0.0, 1.0).cpu() - return m[..., ch:ch + 1].expand(-1, -1, 3).contiguous() + return torch.zeros((*ref.shape[:-1], 1), dtype=torch.float32).expand(-1, -1, 3) + m = packed_mr[..., ch:ch + 1].float().clamp(0.0, 1.0).cpu() + return m.expand(-1, -1, 3) class BakeTextureFromVoxel(IO.ComfyNode): @@ -1375,6 +1412,9 @@ class BakeTextureFromVoxel(IO.ComfyNode): @classmethod def execute(cls, mesh, voxel_colors, texture_size, reference_mesh=None): + reference_faces = _mesh_face_count(reference_mesh if reference_mesh is not None else mesh) + memory_required = max(texture_size * texture_size * 512, reference_faces * 512) + _prepare_gpu_mesh_processing(comfy.model_management.get_torch_device(), memory_required) voxels = voxel_colors coords = voxels.data colors = voxels.voxel_colors @@ -1424,9 +1464,17 @@ class BakeTextureFromVoxel(IO.ComfyNode): black = torch.zeros((1, texture_size, texture_size, 3)) return IO.NodeOutput(black, black, black) # Stack [B,H,W,3]; split packed MR (G=roughness, B=metallic) into grayscale maps. - base_img = torch.stack([t.float().clamp(0.0, 1.0).cpu() for t in out_tex], dim=0) - metallic_img = torch.stack([_mr_channel(m, 2, out_tex[0]) for m in out_mr], dim=0) - roughness_img = torch.stack([_mr_channel(m, 1, out_tex[0]) for m in out_mr], dim=0) + base_maps = [t.float().clamp(0.0, 1.0).cpu() for t in out_tex] + metallic_maps = [_mr_channel(m, 2, out_tex[0]) for m in out_mr] + roughness_maps = [_mr_channel(m, 1, out_tex[0]) for m in out_mr] + if len(base_maps) == 1: + base_img = base_maps[0].unsqueeze(0) + metallic_img = metallic_maps[0].unsqueeze(0) + roughness_img = roughness_maps[0].unsqueeze(0) + else: + base_img = torch.stack(base_maps, dim=0) + metallic_img = torch.stack(metallic_maps, dim=0) + roughness_img = torch.stack(roughness_maps, dim=0) return IO.NodeOutput(base_img, metallic_img, roughness_img) # Single-item path. @@ -1586,7 +1634,7 @@ class ApplyTextureToMesh(IO.ComfyNode): def _chan(img, default): if img is None: return torch.full((B, H, W, 1), float(default)) - t = img.float().clamp(0.0, 1.0).cpu()[..., 0:1] + t = img[..., 0:1].float().clamp(0.0, 1.0).cpu() if int(t.shape[1]) != H or int(t.shape[2]) != W: t = torch.nn.functional.interpolate(t.permute(0, 3, 1, 2), size=(H, W), mode="bilinear", align_corners=False).permute(0, 2, 3, 1) @@ -2157,7 +2205,7 @@ def fill_holes_v2_fn(vertices, faces, max_perimeter=0.03, colors=None, weld_epsi def _process_mesh_batch(mesh, per_item_fn): """Dispatch list/batched/single mesh, extract colors, stack results.""" - mesh = copy.deepcopy(mesh) + mesh = copy.copy(mesh) def process_single(v, f, c, bar): v, f, c = per_item_fn(v, f, c) @@ -2390,6 +2438,8 @@ class RemeshMesh(IO.ComfyNode): # ComfyUI passes meshes on CPU (remesh far faster on GPU); compute on device, return on original. compute_device = comfy.model_management.get_torch_device() + memory_required = max(resolution ** 3 * 64, _mesh_face_count(mesh) * 512) + _prepare_gpu_mesh_processing(compute_device, memory_required) counts = {"in": 0, "out": 0} def _fn(v, f, c): @@ -2421,7 +2471,7 @@ class RemeshMesh(IO.ComfyNode): scale=rs_scale, center=rs_center, colors=cc) v = rv.to(src_device) - f = rf.to(src_device) + f = rf.to(device=src_device, dtype=torch.int32) c = rc.to(src_device) if rc is not None else None counts["out"] += int(f.shape[0]) return v, f, c @@ -2696,6 +2746,7 @@ class UnwrapMesh(IO.ComfyNode): @classmethod def execute(cls, mesh, segmenter, resolution, padding, weld_distance): compute_device = comfy.model_management.get_torch_device() + _prepare_gpu_mesh_processing(compute_device, _mesh_face_count(mesh) * 14 * 1024) seg_device = compute_device if segmenter == "pec" else torch.device("cpu") is_list = isinstance(mesh.vertices, list) @@ -2727,7 +2778,7 @@ class UnwrapMesh(IO.ComfyNode): uvs[:, 1] = 1.0 - uvs[:, 1] # UV y flipped vs trimesh out_v.append(torch.from_numpy(vnp[vmapping]).to(src_device)) - out_f.append(torch.from_numpy(indices).to(device=src_device, dtype=torch.long)) + out_f.append(torch.from_numpy(indices).to(device=src_device, dtype=torch.int32)) out_uv.append(torch.from_numpy(uvs.astype(np.float32)).to(src_device)) if ci is not None: cnp = ci.detach().cpu().numpy() diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py index 57bd2c4cc..1e5f535cc 100644 --- a/comfy_extras/nodes_trellis2.py +++ b/comfy_extras/nodes_trellis2.py @@ -21,6 +21,22 @@ def shape_norm(shape_latent, coords): return SparseTensor(feats=feats, coords=coords) +def _move_sparse_tensor_uncached(tensor, device): + return SparseTensor( + feats=tensor.feats.to(device), + coords=tensor.coords.to(device), + shape=tensor.shape, + scale=tensor._scale, + ) + + +def _sparse_vae_decode_memory(point_count, dtype): + # Last 128-channel stage: feature intermediates plus 27 int32 neighbor indices, + # plus sparse-convolution workspace. + bytes_per_point = 896 * comfy.model_management.dtype_size(dtype) + 27 * 4 + return 2 * 1024 ** 3 + int(point_count) * bytes_per_point + + def infer_batched_coord_layout(coords): if coords.ndim != 2 or coords.shape[1] != 4: raise ValueError(f"Expected Trellis2 coords with shape [N, 4], got {tuple(coords.shape)}") @@ -127,7 +143,11 @@ class VaeDecodeShapeTrellis(IO.ComfyNode): sample_tensor = samples["samples"] device = comfy.model_management.get_torch_device() coords = samples["coords"] - vae.prepare_decode(sample_tensor.shape) + surface_point_estimate = resolution * resolution * 5 // 4 + vae.prepare_decode( + sample_tensor.shape, + memory_required=_sparse_vae_decode_memory(surface_point_estimate, vae.vae_dtype), + ) trellis_vae = vae.first_stage_model coord_counts = samples.get("coord_counts") @@ -166,6 +186,8 @@ class VaeDecodeShapeTrellis(IO.ComfyNode): mesh = Types.MESH(vertices=torch.stack(vert_list), faces=torch.stack(face_list)) else: mesh = pack_variable_mesh_batch(vert_list, face_list) + output_device = comfy.model_management.intermediate_device() + subs = [_move_sparse_tensor_uncached(sub, output_device) for sub in subs] return IO.NodeOutput(mesh, subs) class VaeDecodeTextureTrellis(IO.ComfyNode): @@ -193,7 +215,10 @@ class VaeDecodeTextureTrellis(IO.ComfyNode): sample_tensor = samples["samples"] device = comfy.model_management.get_torch_device() coords = samples["coords"] - vae.prepare_decode(sample_tensor.shape) + vae.prepare_decode( + sample_tensor.shape, + memory_required=_sparse_vae_decode_memory(shape_subdivides[-1].feats.shape[0], vae.vae_dtype), + ) trellis_vae = vae.first_stage_model coord_counts = samples.get("coord_counts") model_frame = samples.get("model_frame", "y_up") @@ -204,6 +229,7 @@ class VaeDecodeTextureTrellis(IO.ComfyNode): samples = samples.to(device) feats = tex_slat_format.process_out(samples) samples = SparseTensor(feats=feats, coords=coords.to(device)) + shape_subdivides = [_move_sparse_tensor_uncached(sub, device) for sub in shape_subdivides] voxel = trellis_vae.decode_tex_slat(samples.to(vae.vae_dtype), shape_subdivides) # Keep all decoded channels. The texture VAE emits 6: base_color (0:3), @@ -239,6 +265,9 @@ class VaeDecodeTextureTrellis(IO.ComfyNode): dim=-1, ) + output_device = comfy.model_management.intermediate_device() + voxel_coords = voxel_coords.to(output_device) + color_feats = color_feats.to(output_device) voxel = Types.VOXEL(voxel_coords, color_feats, tex_resolution) return IO.NodeOutput(voxel)