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)

comfy.model_management.free_memory(

int(memory_required) + comfy.model_management.minimum_inference_memory(),

device,

"""Paint a mesh using nearest-neighbor colors from a sparse voxel field."""

device = comfy.model_management.vae_offload_device()

origin = torch.tensor([-0.5, -0.5, -0.5], device=device)

voxel_size = 1.0 / resolution

voxel_pos = voxel_coords.to(device).float() * voxel_size + origin

verts = mesh.vertices.to(device).squeeze(0)

voxel_colors = voxel_colors.to(device)

voxel_pos_np = voxel_pos.cpu().numpy()

verts_np = verts.cpu().numpy()

tree = cKDTree(voxel_pos_np)

_, nearest_idx_np = tree.query(verts_np, k=1, workers=-1)

nearest_idx = torch.from_numpy(nearest_idx_np).long().to(voxel_colors.device)

v_colors = voxel_colors[nearest_idx]

# Voxel field may carry full PBR; vertex colors use only base_color RGB.

if v_colors.shape[-1] > 3:

v_colors = v_colors[:, :3]

srgb_colors = v_colors.clamp(0, 1)

# to Linear RGB (required for GLTF)

linear_colors = torch.pow(srgb_colors, 2.2)

final_colors = linear_colors.unsqueeze(0)

out_mesh = copy.deepcopy(mesh)

out_mesh.vertex_colors = final_colors

out_mesh = copy.copy(mesh)

vertex_count = mesh.vertices.shape[1]

out_mesh.vertex_colors = mesh.vertices.new_zeros((1, vertex_count, 3))

@classmethod

def define_schema(cls):

return IO.Schema(

node_id="PaintMesh",

display_name="Paint Mesh",

category="3d/mesh",

description="Paints each mesh vertex with its nearest voxel color from the input voxel field.",

inputs=[

IO.Mesh.Input("mesh"),

IO.Voxel.Input("voxel_colors")

],

outputs=[

IO.Mesh.Output("mesh"),

]

)

@classmethod

def execute(cls, mesh, voxel_colors):

voxels = voxel_colors

coords = voxels.data

colors = voxels.voxel_colors

resolution = voxels.resolution

if coords.shape[0] == 0:

return IO.NodeOutput(paint_mesh_default_colors(mesh))

mesh_batch_size = mesh.vertices.shape[0]

if coords.shape[-1] == 4 and mesh_batch_size > 1:

batch_idx = coords[:, 0].long()

voxel_coords = coords[:, 1:]

mesh_batch_size = mesh.vertices.shape[0]

out_verts, out_faces, out_colors = [], [], []

for i in range(mesh_batch_size):

sel = batch_idx == i

item_coords = voxel_coords[sel]

item_colors = colors[sel]

item_vertices, item_faces, *_ = get_mesh_batch_item(mesh, i)

item_mesh = Types.MESH(vertices=item_vertices.unsqueeze(0), faces=item_faces.unsqueeze(0))

if item_coords.shape[0] == 0:

painted = paint_mesh_default_colors(item_mesh)

else:

painted = paint_mesh_with_voxels(item_mesh, item_coords, item_colors, resolution=resolution)

out_verts.append(painted.vertices.squeeze(0))

out_faces.append(painted.faces.squeeze(0))

out_colors.append(painted.vertex_colors.squeeze(0))

out_mesh = pack_variable_mesh_batch(out_verts, out_faces, out_colors)

return IO.NodeOutput(out_mesh)

if coords.shape[-1] == 4:

coords = coords[:, 1:]

out_mesh = paint_mesh_with_voxels(mesh, coords, colors, resolution=resolution)

"""Rasterize the mesh in UV space (tiled point-in-triangle, pure torch). Returns per-texel

dev = comfy.model_management.get_torch_device()

H = W = int(texture_size)

face_idx = torch.zeros((H, W), dtype=torch.long, device=dev)

bary = torch.zeros((H, W, 3), device=dev)

cov = torch.zeros((H, W), dtype=torch.bool, device=dev)

if faces_np.shape[0] == 0:

return face_idx, bary, cov

uvs = torch.from_numpy(np.ascontiguousarray(uvs_np, dtype=np.float32)).to(dev)

faces = torch.from_numpy(np.ascontiguousarray(faces_np).astype(np.int64)).to(dev)

# GL convention: window coord = uv * resolution, coverage tested at texel centre.

tri_uv = (uvs * float(W))[faces] # [F,3,2]

x0, y0 = tri_uv[:, 0, 0], tri_uv[:, 0, 1]

x1, y1 = tri_uv[:, 1, 0], tri_uv[:, 1, 1]

x2, y2 = tri_uv[:, 2, 0], tri_uv[:, 2, 1]

denom = (y1 - y2) * (x0 - x2) + (x2 - x1) * (y0 - y2)

nondegen = denom.abs() > 1e-20

xmin = torch.minimum(torch.minimum(x0, x1), x2).floor().clamp_(0, W - 1).long()

xmax = torch.maximum(torch.maximum(x0, x1), x2).ceil().clamp_(0, W - 1).long()

ymin = torch.minimum(torch.minimum(y0, y1), y2).floor().clamp_(0, H - 1).long()

ymax = torch.maximum(torch.maximum(y0, y1), y2).ceil().clamp_(0, H - 1).long()

# Tile so point-in-triangle only runs over the triangles whose bbox hits the tile.

TILE = 64

eps = 1e-6

for ty in range(0, H, TILE):

ty_end = min(ty + TILE, H)

for tx in range(0, W, TILE):

tx_end = min(tx + TILE, W)

tri_mask = (nondegen & (xmin < tx_end) & (xmax >= tx)

& (ymin < ty_end) & (ymax >= ty))

if not tri_mask.any():

continue

idx = torch.nonzero(tri_mask, as_tuple=True)[0]

ys = torch.arange(ty, ty_end, dtype=torch.float32, device=dev) + 0.5

xs = torch.arange(tx, tx_end, dtype=torch.float32, device=dev) + 0.5

yy, xx = torch.meshgrid(ys, xs, indexing="ij") # [th,tw]

sx0, sy0 = x0[idx][:, None, None], y0[idx][:, None, None]

sx1, sy1 = x1[idx][:, None, None], y1[idx][:, None, None]

sx2, sy2 = x2[idx][:, None, None], y2[idx][:, None, None]

sden = denom[idx][:, None, None]

b0 = ((sy1 - sy2) * (xx - sx2) + (sx2 - sx1) * (yy - sy2)) / sden

b1 = ((sy2 - sy0) * (xx - sx2) + (sx0 - sx2) * (yy - sy2)) / sden

b2 = 1.0 - b0 - b1

inside = (b0 >= -eps) & (b1 >= -eps) & (b2 >= -eps) # [K,th,tw]

if not inside.any():

continue

hit = inside.any(dim=0) # [th,tw]

sel = inside.int().argmax(dim=0) # [th,tw] first covering local tri

bsel = torch.stack([b0.gather(0, sel[None]).squeeze(0),

b1.gather(0, sel[None]).squeeze(0),

b2.gather(0, sel[None]).squeeze(0)], dim=-1) # [th,tw,3]

face_idx[ty:ty_end, tx:tx_end][hit] = idx[sel][hit] # slice is a view → writes through

bary[ty:ty_end, tx:tx_end][hit] = bsel[hit]

cov[ty:ty_end, tx:tx_end] |= hit

"""Interpolate a per-vertex attribute [N,C] into a [H,W,C] map via a rasterized

H, W = mask.shape

out = torch.zeros((H, W, attr_v.shape[1]), device=attr_v.device, dtype=attr_v.dtype)

if mask.any():

vtri = attr_v[faces[face_idx[mask]]] # [K,3,C]

out[mask] = (bary[mask][:, :, None] * vtri).sum(1)

"""Barycentric-interpolate a per-vertex vec3 (world position, or any vec3 e.g. normals)

dev = comfy.model_management.get_torch_device()

H = W = int(texture_size)

if faces_np.shape[0] == 0:

return np.zeros((H, W, 3), dtype=np.float32), np.zeros((H, W), dtype=bool)

face_idx, bary, mask = _rasterize_uv_barycentric(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

"""Normalized trilinear over a SPARSE voxel field (only occupied corners of the 8,

R = int(resolution)

origin = -0.5

voxel_size = 1.0 / R

# Cell-CENTER convention: coord c sits at origin+(c+0.5)*voxel_size (matches

# official grid_sample_3d); the -0.5 puts integer gc on centres so the 8 corners

# bracket the query (omitting it bleeds colour at boundaries/thin features).

gc = (positions.astype(np.float64) - origin) / voxel_size - 0.5

base = np.floor(gc).astype(np.int64)

frac = gc - base

vc = voxel_coords_np.astype(np.int64)

occ_keys = (vc[:, 0] * R + vc[:, 1]) * R + vc[:, 2]

order = np.argsort(occ_keys)

occ_sorted = occ_keys[order]

K = positions.shape[0]

C = color_np.shape[1]

acc = np.zeros((K, C), dtype=np.float64)

wsum = np.zeros((K, 1), dtype=np.float64)

for dx in (0, 1):

wx = frac[:, 0] if dx else 1.0 - frac[:, 0]

for dy in (0, 1):

wy = frac[:, 1] if dy else 1.0 - frac[:, 1]

for dz in (0, 1):

wz = frac[:, 2] if dz else 1.0 - frac[:, 2]

cx = base[:, 0] + dx

cy = base[:, 1] + dy

cz = base[:, 2] + dz

inb = (cx >= 0) & (cx < R) & (cy >= 0) & (cy < R) & (cz >= 0) & (cz < R)

key = (cx * R + cy) * R + cz

ins = np.clip(np.searchsorted(occ_sorted, key), 0, len(occ_sorted) - 1)

matched = inb & (occ_sorted[ins] == key)

idx = order[ins] # garbage where !matched

w = np.where(matched, wx * wy * wz, 0.0)[:, None]

acc += w * color_np[idx] # w=0 cancels garbage rows

wsum += w

ok = wsum[:, 0] > 1e-8

vals = np.zeros((K, C), dtype=np.float64)

vals[ok] = acc[ok] / wsum[ok]

"""GPU port of `_trilinear_sample_sparse`. Returns (vals [K,C] float32, ok [K] bool)."""

dev = comfy.model_management.get_torch_device()

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(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).to(index_dtype)

frac = gc - base.float()

key = (VC[:, 0] * R + VC[:, 1]) * R + VC[:, 2]

skey, order = key.sort()

acc = torch.zeros((K, C), device=dev)

wsum = torch.zeros((K, 1), device=dev)

for dx in (0, 1):

wx = frac[:, 0] if dx else 1.0 - frac[:, 0]

for dy in (0, 1):

wy = frac[:, 1] if dy else 1.0 - frac[:, 1]

for dz in (0, 1):

wz = frac[:, 2] if dz else 1.0 - frac[:, 2]

cx = base[:, 0] + dx

cy = base[:, 1] + dy

cz = base[:, 2] + dz

inb = (cx >= 0) & (cx < R) & (cy >= 0) & (cy < R) & (cz >= 0) & (cz < R)

qk = (cx * R + cy) * R + cz

ins = torch.searchsorted(skey, qk).clamp(max=M - 1)

matched = inb & (skey[ins] == qk)

idx = order[ins] # garbage where !matched

w = torch.where(matched, wx * wy * wz, torch.zeros_like(wx))[:, None]

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)

vals[ok] = acc[ok] / wsum[ok].clamp_min(1e-8)

"""GPU nearest-occupied-voxel lookup via sorted-key grid scan. Returns (vals [K,C]

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

skey, order = key.sort()

def _search(idx, radius):

"""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().to(index_dtype)

n = idx.shape[0]

bd = torch.full((n,), 1e30, device=dev)

bi = torch.zeros(n, dtype=torch.long, device=dev)

fnd = torch.zeros(n, dtype=torch.bool, device=dev)

rng = range(-radius, radius + 1)

for dx in rng:

for dy in rng:

for dz in rng:

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)

match = inb & (skey[ins] == qk)

dd = (((cc.float() + 0.5) / R - 0.5 - Ps) ** 2).sum(1)

upd = match & (dd < bd)

bd = torch.where(upd, dd, bd)

bi = torch.where(upd, order[ins], bi)

fnd |= match

return bi, fnd

def _brute_nearest(idx):

"""Exact nearest occupied voxel for the few grid-scan stragglers, chunked GPU

Ps = P[idx] # [N,3] world

N = Ps.shape[0]

vox_pos = (VC.float() + 0.5) / R - 0.5 # [M,3] voxel centres

best_d = torch.full((N,), 1e30, device=dev)

best_j = torch.zeros(N, dtype=torch.long, device=dev)

# Bound the N×chunk matrix to ~64M elements.

chunk = max(1, (1 << 26) // max(1, N))

for s in range(0, M, chunk):

vc = vox_pos[s:s + chunk] # [B,3]

dd = (Ps[:, None, :] - vc[None, :, :]).pow(2).sum(-1) # [N,B]

md, mj = dd.min(1)

upd = md < best_d

best_d = torch.where(upd, md, best_d)

best_j = torch.where(upd, mj + s, best_j)

return best_j

all_idx = torch.arange(K, device=dev)

best_i = torch.zeros(K, dtype=torch.long, device=dev)

found = torch.zeros(K, dtype=torch.bool, device=dev)

# Pass 1: radius 1 over everything; Pass 2: radius 4 on misses; Pass 3: brute force.

bi1, fnd1 = _search(all_idx, 1)

best_i[all_idx] = bi1

found[all_idx] = fnd1

miss = torch.nonzero(~found, as_tuple=True)[0]

if miss.numel() > 0:

bi2, fnd2 = _search(miss, 4)

best_i[miss] = bi2

found[miss] = fnd2

# Pass 3: stragglers >4 cells from any voxel. A handful → GPU brute force; many

# (coarse mesh, texels far from the voxel shell) → leave unfound for the caller's

# cKDTree, since brute force is O(N·M) and its chunk loop blows up at large N.

miss2 = torch.nonzero(~found, as_tuple=True)[0]

if 0 < miss2.numel() <= _BRUTE_NEAREST_MAX:

best_i[miss2] = _brute_nearest(miss2)

found[miss2] = True

vals = col[best_i]

"""Sample all voxel attribute channels at every masked texel. Returns (H,W,C)

H, W, _ = position_map.shape

color_np = voxel_colors.detach().cpu().numpy().astype(np.float32)

C = color_np.shape[-1]

out = np.zeros((H, W, C), dtype=np.float32)

if not mask.any():

return out

coords_np = voxel_coords.detach().cpu().numpy()

valid_positions = position_map[mask]

def _nearest(query):

# Grid scan + small-N brute tail (on the compute device). Only a large count of far

# stragglers (coarse mesh, or a surface off the voxel shell) is left unfound → resolve

# those with one cKDTree, since GPU brute force is O(N·M) and blows up at large N.

vals, found = _nearest_voxel_sample_gpu(query, coords_np, color_np, resolution)

if not found.all():

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]

return vals

try:

vals, ok = _trilinear_sample_sparse_gpu(valid_positions, coords_np, color_np, resolution)

except Exception as e:

comfy.model_management.raise_non_oom(e) # only fall back on OOM; surface real errors

logging.warning(f"[BakeTextureFromVoxel] GPU trilinear ran out of memory ({e}); falling back to CPU")

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

np.clip(vals, 0.0, 1.0, out=vals)

out[mask] = vals

"""floor(log2(x)) elementwise for int64 x >= 1 (bit-search, no float)."""

r = torch.zeros_like(x)

xx = x.clone()

for s in (32, 16, 8, 4, 2, 1):

sh = xx >> s

m = sh > 0

r = torch.where(m, r + s, r)

xx = torch.where(m, sh, xx)

"""Spread the low 21 bits of v across every 3rd bit (for a 63-bit Morton code)."""

v = v & 0x1fffff

v = (v | (v << 32)) & 0x1f00000000ffff

v = (v | (v << 16)) & 0x1f0000ff0000ff

v = (v | (v << 8)) & 0x100f00f00f00f00f

v = (v | (v << 4)) & 0x10c30c30c30c30c3

v = (v | (v << 2)) & 0x1249249249249249

"""Linear BVH (Karras 2012) over triangle AABBs, pure torch, no external deps

dev = tri.device

T = tri.shape[0]

amin = tri.amin(1)

amax = tri.amax(1)

cent = (amin + amax) * 0.5

lo = cent.amin(0)

hi = cent.amax(0)

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_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.

def delta(i, j):

ok = (j >= 0) & (j < T)

jj = j.clamp(0, T - 1)

x = msort[i] ^ msort[jj]

same = x == 0

cp = torch.where(same, torch.full_like(x, 63), 62 - _msb_int64(x.clamp_min(1)))

xi = i ^ jj

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

dmin = torch.minimum(dplus, dminus)

# range length: exponential probe then binary search

lmax = torch.full_like(I, 2)

while True:

cond = delta(I, I + lmax * direction) > dmin

if not bool(cond.any()):

break

lmax = torch.where(cond, lmax * 2, lmax)

if int(lmax.max()) > 2 * T:

break

l = torch.zeros_like(I)

t = lmax.clone()

while True:

t = t // 2

if int(t.max()) == 0:

break

cond = delta(I, I + (l + t) * direction) > dmin

l = torch.where(cond, l + t, l)

j = I + l * direction

first = torch.minimum(I, j)

last = torch.maximum(I, j)

# split position: binary search on delta within [first, last]

dnode = delta(first, last)

s = torch.zeros_like(I)

div = torch.full_like(I, 2)

rng = last - first

while True:

step = (rng + div - 1) // div

cond = delta(first, (first + s + step).clamp(max=T - 1)) > dnode

s = torch.where(cond, s + step, s)

if int(step.max()) <= 1:

cond1 = delta(first, (first + s + 1).clamp(max=T - 1)) > dnode

s = torch.where(cond1, s + 1, s)

break

div = div * 2

gamma = first + s

LEAF = T

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

need = ~setm[:T - 1]

if not bool(need.any()):

break

idx = torch.nonzero(need, as_tuple=True)[0]

ii = idx[setm[left[idx]] & setm[right[idx]]]

if ii.numel() == 0:

break

nmin[ii] = torch.minimum(nmin[left[ii]], nmin[right[ii]])

nmax[ii] = torch.maximum(nmax[left[ii]], nmax[right[ii]])

setm[ii] = True

"""Exact closest surface point per query via per-query BVH stack traversal

dev = Q.device

N = Q.shape[0]

LEAF = bvh['LEAF']

nmin = bvh['nmin']

nmax = bvh['nmax']

left = bvh['left']

right = bvh['right']

order = bvh['order']

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)

bestp = Q.clone()

bestf = torch.full((N,), -1, dtype=torch.long, device=dev)

active = torch.arange(N, device=dev)

overflow = 0

def aabb_d2(node, q):

d = (nmin[node] - q).clamp_min(0) + (q - nmax[node]).clamp_min(0)

return (d * d).sum(-1)

while active.numel() > 0:

a = active

qa = Q[a]

node = stack[a, sp[a] - 1]

sp[a] = sp[a] - 1

within = aabb_d2(node, qa) < best[a]

isleaf = node >= LEAF

lv = within & isleaf

if bool(lv.any()):

ga = a[lv]

fidx = order[node[lv] - LEAF] # triangle index of each leaf

tt = tri[fidx]

cp, d2 = _point_tri_closest(qa[lv], tt)

upd = d2 < best[ga]

gu = ga[upd]

best[gu] = d2[upd]

bestp[gu] = cp[upd]

bestf[gu] = fidx[upd].long()

iv = within & ~isleaf

if bool(iv.any()):

gi = a[iv]

qi = qa[iv]

lc = left[node[iv]]

rc = right[node[iv]]

dl = aabb_d2(lc, qi)

dr = aabb_d2(rc, qi)

near = torch.where(dl <= dr, lc, rc)

far = torch.where(dl <= dr, rc, lc)

s0 = sp[gi]

stack[gi, s0.clamp(max=max_stack - 1)] = far

sp[gi] = (s0 + 1).clamp(max=max_stack)

s1 = sp[gi]

overflow += int((s1 >= max_stack).sum())

stack[gi, s1.clamp(max=max_stack - 1)] = near

sp[gi] = (s1 + 1).clamp(max=max_stack)

active = a[sp[a] > 0]

if overflow:

logging.warning(f"[back-project] BVH stack overflow on {overflow} pushes "

f"(max_stack={max_stack}); a few texels may be slightly off — "

f"raise max_stack if this is large.")

if return_face:

return bestp, bestf

"""Snap covered texels onto the reference mesh's true surface (pure-torch BVH, no

if not mask.any():

return position_map

dev = comfy.model_management.get_torch_device()

rv = ref_v.detach().to(dev).float()

rf = ref_f.detach().to(device=dev, dtype=torch.int32)

tri = rv[rf]

bvh = _build_triangle_bvh(tri)

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

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.

pos_lo = position_map[::stride, ::stride]

mask_lo = mask[::stride, ::stride]

Hl, Wl = mask_lo.shape

corr_lo = np.zeros((Hl, Wl, 3), dtype=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(

torch.from_numpy(np.ascontiguousarray(corr_lo)).permute(2, 0, 1)[None].to(dev),

size=(H, W), mode="bilinear", align_corners=False,

)[0].permute(1, 2, 0).cpu().numpy()

out = position_map.copy()

out[mask] = position_map[mask] + corr[mask]

"""Möller-Trumbore any-hit per (ray, triangle) pair, double-sided. Returns bool [N]."""

a, b, c = tri[:, 0], tri[:, 1], tri[:, 2]

e1, e2 = b - a, c - a

p = torch.cross(d, e2, dim=-1)

det = (e1 * p).sum(-1)

inv = 1.0 / torch.where(det.abs() < 1e-20, torch.full_like(det, 1e-20), det)

tvec = o - a

u = (tvec * p).sum(-1) * inv

q = torch.cross(tvec, e1, dim=-1)

v = (d * q).sum(-1) * inv

t = (e2 * q).sum(-1) * inv

"""Any-hit ray test over the BVH (slab cull + Möller-Trumbore), pure torch. Returns bool

dev = orig.device

N = orig.shape[0]

LEAF = bvh['LEAF']

nmin, nmax = bvh['nmin'], bvh['nmax']

left, right, order = bvh['left'], bvh['right'], bvh['order']

inv = 1.0 / torch.where(dirs.abs() < 1e-20, torch.full_like(dirs, 1e-20), dirs)

tmaxN = tmax if torch.is_tensor(tmax) else torch.full((N,), float(tmax), device=dev) # per-ray far bound

hit = torch.zeros(N, dtype=torch.bool, device=dev)

# int32 stack: node indices fit in 31 bits and this [N, max_stack] array dominates memory.

stack = torch.full((N, max_stack), -1, dtype=torch.int32, device=dev)

sp = torch.ones(N, dtype=torch.long, device=dev)

stack[:, 0] = 0

active = torch.arange(N, device=dev)

def slab(node, o, i, tmx):

t1 = (nmin[node] - o) * i

t2 = (nmax[node] - o) * i

tnear = torch.minimum(t1, t2).amax(-1)

tfar = torch.maximum(t1, t2).amin(-1)

return (tfar >= tnear.clamp_min(tmin)) & (tnear <= tmx) & (tfar >= tmin)

while active.numel() > 0:

a = active

node = stack[a, sp[a] - 1]

sp[a] = sp[a] - 1

within = slab(node, orig[a], inv[a], tmaxN[a])

isleaf = node >= LEAF

lv = within & isleaf

if bool(lv.any()):

ga = a[lv]

tt = tri[order[node[lv] - LEAF]]

h = _ray_tri_hit(orig[ga], dirs[ga], tt, tmin, tmaxN[ga])

hit[ga[h]] = True

iv = within & ~isleaf

if bool(iv.any()):

gi = a[iv]

s0 = sp[gi]