Lower peak trellis workflow vram and ram usage. (#16054)

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)
 

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

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.

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:

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)

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()

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)