wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-11 08:57:22 +08:00
Code Issues Actions 17 Packages Projects Releases Wiki Activity

Refactor/optimize UV unwrap

Browse Source
This commit is contained in:
kijai 2026-07-03 20:23:26 +03:00
parent 3b34e177cb
commit 5067ac461e
4 changed files with 1196 additions and 986 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1201
comfy_extras/mesh3d/uv_unwrap/pack.py
View File

File diff suppressed because it is too large Load Diff

316
comfy_extras/mesh3d/uv_unwrap/parameterize.py
View File

@ -12,6 +12,8 @@ from torch import Tensor
from . import mesh as _mesh from . import mesh as _mesh
LSCM_BATCH_MAX_VERTS = 256 # charts above this solve per-chart sparse (lscm_chart)
def solve_least_squares(A: sp.csr_matrix, b: np.ndarray) -> np.ndarray: def solve_least_squares(A: sp.csr_matrix, b: np.ndarray) -> np.ndarray:
"""Solve ||Ax - b||^2 by factorizing AtA.""" """Solve ||Ax - b||^2 by factorizing AtA."""
@ -99,54 +101,258 @@ def _ortho_project(verts_3d: np.ndarray) -> np.ndarray:
return np.stack([verts_3d @ t, verts_3d @ b], axis=1) return np.stack([verts_3d @ t, verts_3d @ b], axis=1)
def _stretch_metrics(verts_3d: np.ndarray, uvs: np.ndarray, faces: np.ndarray) -> Tuple[float, float, int, int]: def ortho_project_concat(verts: np.ndarray, chart_of_vert: np.ndarray, n_charts: int) -> np.ndarray:
"""Sander's stretch metric. Returns (rms, max, n_flipped, n_zero_area).""" """_ortho_project for every chart at once over concatenated per-chart vertices."""
p = verts_3d[faces] cnt = np.bincount(chart_of_vert, minlength=n_charts).clip(min=1).astype(np.float64)
cen = np.stack([np.bincount(chart_of_vert, weights=verts[:, i], minlength=n_charts)
for i in range(3)], axis=1) / cnt[:, None]
d = verts - cen[chart_of_vert]
cov = np.zeros((n_charts, 3, 3), dtype=np.float64)
for i in range(3):
for j in range(i, 3):
s = np.bincount(chart_of_vert, weights=d[:, i] * d[:, j], minlength=n_charts)
cov[:, i, j] = s
cov[:, j, i] = s
_w, ev = np.linalg.eigh(cov)
normal = ev[:, :, 0]
t = np.eye(3, dtype=np.float64)[np.argmin(np.abs(normal), axis=1)]
t = t - normal * (normal * t).sum(axis=1, keepdims=True)
t /= np.linalg.norm(t, axis=1, keepdims=True).clip(min=1e-20)
b = np.cross(normal, t)
tt, bb = t[chart_of_vert], b[chart_of_vert]
return np.stack([(verts * tt).sum(1), (verts * bb).sum(1)], axis=1)
def stretch_metrics_concat(
verts: np.ndarray, uvs: np.ndarray, faces: np.ndarray,
chart_of_face: np.ndarray, n_charts: int,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Per-chart Sander stretch metrics (rms, max, n_flipped, n_zero_area); rms/max inf where undefined."""
p = verts[faces]
t = uvs[faces] t = uvs[faces]
parametric_area = 0.5 * ( pa_signed = 0.5 * (
(t[:, 1, 1] - t[:, 0, 1]) * (t[:, 2, 0] - t[:, 0, 0]) (t[:, 1, 1] - t[:, 0, 1]) * (t[:, 2, 0] - t[:, 0, 0])
- (t[:, 2, 1] - t[:, 0, 1]) * (t[:, 1, 0] - t[:, 0, 0]) - (t[:, 2, 1] - t[:, 0, 1]) * (t[:, 1, 0] - t[:, 0, 0]))
) n_flip = np.bincount(chart_of_face[pa_signed < -1e-12], minlength=n_charts)
n_flipped = int((parametric_area < -1e-12).sum()) n_zero = np.bincount(chart_of_face[np.abs(pa_signed) < 1e-12], minlength=n_charts)
n_zero = int((np.abs(parametric_area) < 1e-12).sum()) pa = np.abs(pa_signed).clip(min=1e-20)
pa = np.abs(parametric_area).clip(min=1e-20) ga = 0.5 * np.linalg.norm(np.cross(p[:, 1] - p[:, 0], p[:, 2] - p[:, 0]), axis=1)
geom_area = 0.5 * np.linalg.norm( keep = (ga > 1e-12) & (np.abs(pa_signed) > 1e-12)
np.cross(p[:, 1] - p[:, 0], p[:, 2] - p[:, 0]), axis=1 t1, s1 = t[:, 0, 0], t[:, 0, 1]
) t2, s2 = t[:, 1, 0], t[:, 1, 1]
keep = (geom_area > 1e-12) & (np.abs(parametric_area) > 1e-12) t3, s3 = t[:, 2, 0], t[:, 2, 1]
if not keep.any():
return float("inf"), float("inf"), n_flipped, n_zero
t1 = t[:, 0, 0]
s1 = t[:, 0, 1]
t2 = t[:, 1, 0]
s2 = t[:, 1, 1]
t3 = t[:, 2, 0]
s3 = t[:, 2, 1]
inv_2pa = 1.0 / (2.0 * pa) inv_2pa = 1.0 / (2.0 * pa)
Ss = ( Ss = (p[:, 0] * (t2 - t3)[:, None] + p[:, 1] * (t3 - t1)[:, None]
p[:, 0] * (t2 - t3)[:, None] + p[:, 2] * (t1 - t2)[:, None]) * inv_2pa[:, None]
+ p[:, 1] * (t3 - t1)[:, None] St = (p[:, 0] * (s3 - s2)[:, None] + p[:, 1] * (s1 - s3)[:, None]
+ p[:, 2] * (t1 - t2)[:, None] + p[:, 2] * (s2 - s1)[:, None]) * inv_2pa[:, None]
) * inv_2pa[:, None]
St = (
p[:, 0] * (s3 - s2)[:, None]
+ p[:, 1] * (s1 - s3)[:, None]
+ p[:, 2] * (s2 - s1)[:, None]
) * inv_2pa[:, None]
a = (Ss * Ss).sum(axis=1) a = (Ss * Ss).sum(axis=1)
bb = (Ss * St).sum(axis=1) bb = (Ss * St).sum(axis=1)
c = (St * St).sum(axis=1) c = (St * St).sum(axis=1)
sigma2_sq = 0.5 * (a + c + np.sqrt(np.maximum(0.0, (a - c) ** 2 + 4 * bb ** 2))) sigma2_sq = 0.5 * (a + c + np.sqrt(np.maximum(0.0, (a - c) ** 2 + 4 * bb ** 2)))
rms_sq = (a + c) * 0.5 rms_sq = (a + c) * 0.5
rms_stretch_sq_sum = float((rms_sq[keep] * geom_area[keep]).sum()) cf = chart_of_face[keep]
total_geom = float(geom_area[keep].sum()) tg = np.bincount(cf, weights=ga[keep], minlength=n_charts)
total_param = float(pa[keep].sum()) tp = np.bincount(cf, weights=pa[keep], minlength=n_charts)
if total_geom <= 0.0: rs = np.bincount(cf, weights=(rms_sq * ga)[keep], minlength=n_charts)
return float("inf"), float("inf"), n_flipped, n_zero smax = np.zeros(n_charts, dtype=np.float64)
norm_factor = np.sqrt(total_param / total_geom) np.maximum.at(smax, cf, sigma2_sq[keep])
rms_stretch = float(np.sqrt(rms_stretch_sq_sum / total_geom)) * norm_factor ok = tg > 0.0
max_stretch = float(np.sqrt(sigma2_sq[keep].max())) * norm_factor tg_safe = np.where(ok, tg, 1.0)
return rms_stretch, max_stretch, n_flipped, n_zero norm = np.sqrt(tp / tg_safe)
rms = np.where(ok, np.sqrt(rs / tg_safe) * norm, np.inf)
mx = np.where(ok, np.sqrt(smax) * norm, np.inf)
return rms, mx, n_flip, n_zero
def _segment_argmax(vals: np.ndarray, seg: np.ndarray, n: int) -> np.ndarray:
"""Index of the (first) max element per segment; -1 for empty segments."""
amax = np.full(n, -np.inf)
np.maximum.at(amax, seg, vals)
hit = vals == amax[seg]
out = np.full(n, np.iinfo(np.int64).max, dtype=np.int64)
np.minimum.at(out, seg[hit], np.nonzero(hit)[0])
return np.where(out == np.iinfo(np.int64).max, -1, out)
def lscm_charts_batch(
verts: np.ndarray, # (sumV, 3) float64, per-chart concatenated
uv_pins: np.ndarray, # (sumV, 2) float64, ortho UVs (pin values + fallback)
faces_gl: np.ndarray, # (sumF, 3) global-local ids into verts
face_pos: np.ndarray, # (sumF,) row index of each face within its chart
chart_of_face: np.ndarray, # (sumF,)
chart_of_vert: np.ndarray, # (sumV,)
vert_offsets: np.ndarray, # (n_charts+1,)
chart_ids: np.ndarray, # charts to solve (each with >=3 verts, >=1 face)
n_charts: int,
max_bucket_verts: int = LSCM_BATCH_MAX_VERTS,
device: "torch.device | None" = None,
) -> dict:
"""Batched dense ABF/LSCM; returns {chart_id: (Vc, 2) float32}. Charts larger than
max_bucket_verts are left out (the caller solves those sparse)."""
out: dict = {}
if chart_ids.size == 0:
return out
sel = np.zeros(n_charts, dtype=bool)
sel[chart_ids] = True
vcounts = np.diff(vert_offsets)
# ABF coefficients for all selected faces in one shot
fmask = sel[chart_of_face]
f_ids = np.nonzero(fmask)[0]
abf_ids, abf_cos, abf_sin, abf_valid = _abf_face_coefficients(verts, faces_gl[f_ids])
# farthest-point pin pair per chart (two passes)
vmask = sel[chart_of_vert]
v_ids = np.nonzero(vmask)[0]
cv = chart_of_vert[v_ids]
first = vert_offsets[:-1]
d0 = ((verts[v_ids] - verts[first[cv]]) ** 2).sum(1)
pin_a = _segment_argmax(d0, cv, n_charts) # global vert index (into v_ids space)
pin_a = np.where(pin_a >= 0, v_ids[pin_a.clip(min=0)], -1)
d1 = ((verts[v_ids] - verts[pin_a.clip(min=0)[cv]]) ** 2).sum(1)
pin_b = _segment_argmax(d1, cv, n_charts)
pin_b = np.where(pin_b >= 0, v_ids[pin_b.clip(min=0)], -1)
# degenerate (all verts coincide): any distinct vert within the chart (Vc >= 3 guaranteed)
alt = np.where(pin_a == first, first + 1, first)
pin_b = np.where(pin_a == pin_b, alt, pin_b)
fcounts = np.bincount(chart_of_face[f_ids], minlength=n_charts)
# size-sorted chunks padded to their own max, bounded by an element budget so one
# face-heavy chart can't inflate a whole chunk
small = chart_ids[vcounts[chart_ids] <= max_bucket_verts]
sorted_ids = small[np.argsort(vcounts[small], kind="stable")]
budget = (96 << 20) // 8 # float64 elements in a chunk's A
chunks = []
cs = 0
fmax_r = vmax_r = 0
for idx in range(sorted_ids.size):
c2 = sorted_ids[idx]
fm2 = max(fmax_r, int(fcounts[c2]))
vm2 = max(vmax_r, int(vcounts[c2]))
nb = idx - cs + 1
if nb > 1 and (nb > 128 or nb * 4 * fm2 * vm2 > budget):
chunks.append((cs, idx))
cs = idx
fmax_r, vmax_r = int(fcounts[c2]), int(vcounts[c2])
else:
fmax_r, vmax_r = fm2, vm2
if sorted_ids.size:
chunks.append((cs, sorted_ids.size))
for s, e in chunks:
cids = sorted_ids[s:e]
B = cids.size
Vmax = int(vcounts[cids].max())
Fmax = int(fcounts[cids].max())
N = 2 * Vmax
R = 2 * Fmax
compact = np.full(n_charts, -1, dtype=np.int64)
compact[cids] = np.arange(B)
fm = compact[chart_of_face[f_ids]] >= 0
fi = f_ids[fm] # face rows for this chunk
bi = compact[chart_of_face[fi]] # chart slot per face
frow = face_pos[fi]
v0 = vert_offsets[chart_of_face[fi]] # local id = global-local - v0
am = fm.nonzero()[0] # index into abf_* arrays
pieces_i: list = []
pieces_v: list = []
def scatter(rows, cols, vals, bsel):
pieces_i.append((bsel * R + rows) * N + cols)
pieces_v.append(vals)
val = abf_valid[am]
ii = am[val]
ids = abf_ids[ii] - v0[val, None] # local vert ids, reordered
cosf, sinf = abf_cos[ii], abf_sin[ii]
rr, bsel = frow[val] * 2, bi[val]
ones = np.ones(ii.size)
for cc2, vv in ((ids[:, 0], cosf - 1.0), (ids[:, 0] + Vmax, -sinf),
(ids[:, 1], -cosf), (ids[:, 1] + Vmax, sinf), (ids[:, 2], ones)):
scatter(rr, cc2, vv, bsel)
for cc2, vv in ((ids[:, 0], sinf), (ids[:, 0] + Vmax, cosf - 1.0),
(ids[:, 1], -sinf), (ids[:, 1] + Vmax, -cosf), (ids[:, 2] + Vmax, ones)):
scatter(rr + 1, cc2, vv, bsel)
inv = ~val
if inv.any():
jj = fi[inv]
tri2d = _triangle_local_2d(verts, faces_gl[jj])
twice = tri2d[:, 1, 0] * tri2d[:, 2, 1] - tri2d[:, 1, 1] * tri2d[:, 2, 0]
w = 1.0 / np.sqrt(2.0 * np.abs(twice).clip(min=1e-20))
rr2, bs2 = frow[inv] * 2, bi[inv]
lids = faces_gl[jj] - v0[inv, None]
for j in range(3):
jp1, jp2 = (j + 1) % 3, (j + 2) % 3
aj = (tri2d[:, jp1, 0] - tri2d[:, jp2, 0]) * w
bj = (tri2d[:, jp1, 1] - tri2d[:, jp2, 1]) * w
vc2 = lids[:, j]
scatter(rr2, vc2, aj, bs2)
scatter(rr2, vc2 + Vmax, -bj, bs2)
scatter(rr2 + 1, vc2, bj, bs2)
scatter(rr2 + 1, vc2 + Vmax, aj, bs2)
flat = np.concatenate(pieces_i)
A = np.bincount(flat, weights=np.concatenate(pieces_v),
minlength=B * R * N).reshape(B, R, N)
# pins: move their columns to the RHS, then constrain via identity rows
voff = vert_offsets[cids]
pa_l = pin_a[cids] - voff
pb_l = pin_b[cids] - voff
pin_cols = np.stack([pa_l, pb_l, pa_l + Vmax, pb_l + Vmax], 1) # (B,4)
pin_vals = np.stack([uv_pins[pin_a[cids], 0], uv_pins[pin_b[cids], 0],
uv_pins[pin_a[cids], 1], uv_pins[pin_b[cids], 1]], 1)
rhs = np.zeros((B, R), dtype=np.float64)
barange = np.arange(B)
for k in range(4):
rhs -= A[barange, :, pin_cols[:, k]] * pin_vals[:, k, None]
A[barange, :, pin_cols[:, k]] = 0.0
# constrained columns: the 4 pins + padding beyond each chart's vert count
vcs = vcounts[cids]
padm = np.arange(Vmax)[None, :] >= vcs[:, None]
con = np.concatenate([padm, padm], axis=1) # (B,N)
np.put_along_axis(con, pin_cols, True, axis=1)
cval = np.zeros((B, N), dtype=np.float64)
np.put_along_axis(cval, pin_cols, pin_vals, axis=1)
# normal equations + batched solve; the fp64 dense algebra goes to the GPU when available
use_gpu = device is not None and device.type == "cuda"
if use_gpu:
A_t = torch.from_numpy(A).to(device)
At = A_t.transpose(1, 2)
AtA = At @ A_t
Atb = (At @ torch.from_numpy(rhs).to(device).unsqueeze(2)).squeeze(2)
con_t = torch.from_numpy(con).to(device)
free2 = (~con_t[:, :, None]) & (~con_t[:, None, :])
AtA = AtA * free2
diag = torch.diagonal(AtA, dim1=1, dim2=2)
# median (not max) positive diagonal: a degenerate face's ~1e19 squared row
# weight would blow a max-scaled eps past the unit ABF rows
dpos = torch.where(diag > 0, diag, torch.full_like(diag, float("nan")))
dsc = 1e-12 * torch.nan_to_num(dpos.nanmedian(dim=1).values, nan=1e-8).clamp_min(1e-20)
diag += torch.where(con_t, torch.ones_like(diag), dsc[:, None].expand_as(diag))
Atb = torch.where(con_t, torch.from_numpy(cval).to(device), Atb)
x = torch.linalg.solve(AtA, Atb).cpu().numpy()
else:
At = A.transpose(0, 2, 1)
AtA = At @ A # batched BLAS dgemm
Atb = (At @ rhs[:, :, None])[:, :, 0]
AtA *= (~con[:, :, None]) & (~con[:, None, :])
dg = AtA.reshape(B, -1)[:, ::N + 1]
# median positive diagonal (see GPU branch): robust to degenerate-face weights
dpos = np.where(dg > 0, dg, np.nan)
with np.errstate(all="ignore"):
dsc = 1e-12 * np.nan_to_num(np.nanmedian(dpos, axis=1), nan=1e-8).clip(min=1e-20)
dg += np.where(con, 1.0, dsc[:, None])
Atb2 = np.where(con, cval, Atb)
x = np.linalg.solve(AtA, Atb2)
for i2, c2 in enumerate(cids):
vc3 = int(vcs[i2])
out[int(c2)] = np.stack([x[i2, :vc3], x[i2, Vmax:Vmax + vc3]], 1).astype(np.float32)
return out
def _uv_boundary_self_intersects( def _uv_boundary_self_intersects(
@ -181,36 +387,6 @@ def _uv_boundary_self_intersects(
return False return False
def parametrize_chart(
local_verts: Tensor, local_faces: Tensor, local_face_face: Tensor
) -> Tensor:
"""Parameterize one chart: ortho first, ABF/LSCM fallback; charts <=5 faces stay ortho."""
verts_np = local_verts.detach().cpu().numpy().astype(np.float64)
faces_np = local_faces.detach().cpu().numpy().astype(np.int64)
if verts_np.shape[0] < 3 or faces_np.shape[0] == 0:
return torch.zeros((verts_np.shape[0], 2), dtype=torch.float32, device=local_verts.device)
ortho = _ortho_project(verts_np)
n_faces = faces_np.shape[0]
if n_faces <= 5:
return torch.from_numpy(ortho.astype(np.float32)).to(local_verts.device)
rms, mx, n_flip, n_zero = _stretch_metrics(verts_np, ortho, faces_np)
flip_ok = n_flip == 0 or n_flip == n_faces
if flip_ok and n_zero == 0 and rms <= 1.5 and mx <= 2.0:
ff_np = local_face_face.detach().cpu().numpy().astype(np.int64)
if not _uv_boundary_self_intersects(ortho, faces_np, ff_np):
return torch.from_numpy(ortho.astype(np.float32)).to(local_verts.device)
uvs_t = lscm_chart(local_verts, local_faces, local_face_face, pin_positions=ortho)
# Collapsed UV island (aspect > 100:1) blows up packing scale; fall back to ortho.
uvs_np = uvs_t.detach().cpu().numpy()
bbox = uvs_np.max(axis=0) - uvs_np.min(axis=0)
bbox_max = float(max(bbox[0], bbox[1], 1e-12))
bbox_min = float(max(min(bbox[0], bbox[1]), 1e-12))
if bbox_max / bbox_min > 100.0:
return torch.from_numpy(ortho.astype(np.float32)).to(local_verts.device)
return uvs_t
def _abf_face_coefficients( def _abf_face_coefficients(
verts_3d: np.ndarray, faces: np.ndarray verts_3d: np.ndarray, faces: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:

499
comfy_extras/mesh3d/uv_unwrap/segment.py
View File

@ -1,22 +1,11 @@
"""Adaptive cost-grow chart segmentation (CPU); numba optional, numpy path is nd-only.""" """Adaptive cost-grow chart segmentation (vectorized torch, CPU or GPU)."""
from __future__ import annotations from __future__ import annotations
from typing import List, Tuple from typing import Tuple
import numpy as np
import torch import torch
from torch import Tensor from torch import Tensor
from tqdm import tqdm
try:
from numba import njit
_HAVE_NUMBA = True
except ImportError:
_HAVE_NUMBA = False
def njit(*args, **kwargs): # noqa: ARG001
def deco(fn):
return fn
return deco if not args else args[0]
from .mesh import MeshData, face_edge_lengths from .mesh import MeshData, face_edge_lengths
@ -28,129 +17,63 @@ DEFAULT_MAX_COST = 2.0
NORMAL_DEVIATION_HARD_CUTOFF = 0.707 # ~75° NORMAL_DEVIATION_HARD_CUTOFF = 0.707 # ~75°
@njit(cache=True, fastmath=False) def _grow_iter(face_chart, frontier, ff, fn, fa, fel, basis, nsum, area, perim, K,
def _cost_grow_iter_jit( nd_cutoff, tau, w_nd, w_round, w_straight):
face_chart: np.ndarray, face_face: np.ndarray, face_normal: np.ndarray, """One grow pass: each frontier face joins its lowest-cost adjacent chart if cost <= tau;
face_area: np.ndarray, face_edge_len: np.ndarray, returns the number of faces assigned."""
chart_basis: np.ndarray, chart_normal_sum: np.ndarray, u = frontier.nonzero(as_tuple=True)[0]
chart_area: np.ndarray, chart_perim: np.ndarray, if u.numel() == 0:
nd_cutoff: float, max_cost: float, return 0
w_nd: float, w_round: float, w_straight: float, nb = ff[u] # (U,3) neighbor face ids
): nbc = torch.where(nb >= 0, face_chart[nb.clamp_min(0)], nb.new_full((), -1))
"""One grow iter: each unassigned face joins its lowest-cost adjacent chart if cost<max_cost.""" valid = nbc >= 0
F = face_chart.shape[0] d = (fn[u][:, None, :] * basis[nbc.clamp_min(0)]).sum(-1)
best_chart_per_face = np.full(F, -1, dtype=np.int64) nd = (1.0 - d).clamp(0.0, 1.0)
best_cost_per_face = np.full(F, np.inf, dtype=np.float32) valid &= nd < nd_cutoff
el = fel[u] # (U,3)
# l_in per candidate chart j: edge k counts if its (assigned) neighbor is in chart j
inm = (nbc[:, :, None] == nbc[:, None, :]) & valid[:, None, :]
l_in = (el[:, None, :] * inm).sum(-1) # (U,3)
tot = el.sum(-1, keepdim=True)
l_out = tot - l_in
ca = area[nbc.clamp_min(0)]
cp = perim[nbc.clamp_min(0)]
new_perim = cp - l_in + l_out
new_r = new_perim * new_perim / (ca + fa[u][:, None]).clamp_min(1e-20)
round_cost = torch.where((cp <= 1e-20) | (ca <= 1e-20) | (new_r <= 1e-20),
torch.zeros_like(new_r),
1.0 - (cp * cp / ca.clamp_min(1e-20)) / new_r.clamp_min(1e-20))
straight_cost = ((l_out - l_in) / tot.clamp_min(1e-20)).clamp(max=0.0)
cost = w_nd * nd + w_round * round_cost + w_straight * straight_cost
cost = torch.where(valid, cost, cost.new_full((), float("inf")))
best_cost, best_j = cost.min(1)
acc = best_cost <= tau
n_acc = int(acc.sum())
if n_acc == 0:
return 0
for f in range(F): f_acc = u[acc]
if face_chart[f] != -1: c_acc = nbc.gather(1, best_j[:, None]).squeeze(1)[acc]
continue nbc_old = nbc[acc] # neighbor charts before this commit
nx = face_normal[f, 0] face_chart[f_acc] = c_acc
ny = face_normal[f, 1] nb_acc = nb[acc]
nz = face_normal[f, 2] nbs_acc = nb_acc.clamp_min(0)
af = face_area[f] nbc_post = torch.where(nb_acc >= 0, face_chart[nbs_acc], nb_acc.new_full((), -1))
for e0 in range(3): # frontier update: committed faces leave; their still-unassigned neighbors enter
nb0 = face_face[f, e0] frontier[f_acc] = False
if nb0 < 0: grow_nb = nbs_acc[(nb_acc >= 0) & (nbc_post < 0)]
continue frontier[grow_nb] = True
c = face_chart[nb0] el_acc = el[acc]
if c < 0: cx = c_acc[:, None]
continue dper = torch.where(nbc_old == cx, -el_acc, # was member: edge turns interior
d = (nx * chart_basis[c, 0] + ny * chart_basis[c, 1] + nz * chart_basis[c, 2]) torch.where(nbc_post == cx, torch.zeros_like(el_acc), # co-committer
nd = np.float32(1.0) - d el_acc)).sum(1) # boundary / other chart
if nd > np.float32(1.0): perim.scatter_add_(0, c_acc, dper)
nd = np.float32(1.0) area.scatter_add_(0, c_acc, fa[f_acc])
if nd < np.float32(0.0): nsum.index_add_(0, c_acc, fn[f_acc] * fa[f_acc, None])
nd = np.float32(0.0) nl = nsum[:K].norm(dim=1, keepdim=True)
if nd >= nd_cutoff: basis[:K] = torch.where(nl > 1e-20, nsum[:K] / nl.clamp_min(1e-20), basis[:K])
continue return n_acc
l_in = np.float32(0.0)
l_out = np.float32(0.0)
for e1 in range(3):
nb1 = face_face[f, e1]
el = face_edge_len[f, e1]
if nb1 < 0:
l_out += el
elif face_chart[nb1] == c:
l_in += el
else:
l_out += el
ca = chart_area[c]
cp = chart_perim[c]
new_perim = cp - l_in + l_out
new_area = ca + af
if cp <= np.float32(1e-20) or ca <= np.float32(1e-20):
round_cost = np.float32(0.0)
else:
old_r = (cp * cp) / ca
new_r = (new_perim * new_perim) / new_area
if new_r <= np.float32(1e-20):
round_cost = np.float32(0.0)
else:
round_cost = np.float32(1.0) - old_r / new_r
denom = l_out + l_in
if denom <= np.float32(1e-20):
straight_cost = np.float32(0.0)
else:
ratio = (l_out - l_in) / denom
if ratio < np.float32(0.0):
straight_cost = ratio
else:
straight_cost = np.float32(0.0)
cost = (w_nd * nd + w_round * round_cost + w_straight * straight_cost)
if cost < best_cost_per_face[f]:
best_cost_per_face[f] = cost
best_chart_per_face[f] = c
n_assigned = 0
for f in range(F):
if face_chart[f] != -1:
continue
if best_chart_per_face[f] < 0:
continue
if best_cost_per_face[f] > max_cost:
continue
c = best_chart_per_face[f]
l_in = np.float32(0.0)
l_out = np.float32(0.0)
for e1 in range(3):
nb1 = face_face[f, e1]
el = face_edge_len[f, e1]
if nb1 < 0:
l_out += el
elif face_chart[nb1] == c:
l_in += el
else:
l_out += el
af = face_area[f]
face_chart[f] = c
chart_normal_sum[c, 0] += face_normal[f, 0] * af
chart_normal_sum[c, 1] += face_normal[f, 1] * af
chart_normal_sum[c, 2] += face_normal[f, 2] * af
chart_area[c] += af
chart_perim[c] = chart_perim[c] - l_in + l_out
nx = chart_normal_sum[c, 0]
ny = chart_normal_sum[c, 1]
nz = chart_normal_sum[c, 2]
nlen = np.sqrt(nx * nx + ny * ny + nz * nz)
if nlen > np.float32(1e-20):
chart_basis[c, 0] = nx / nlen
chart_basis[c, 1] = ny / nlen
chart_basis[c, 2] = nz / nlen
n_assigned += 1
return n_assigned
def _renumber(face_chart: np.ndarray, device) -> Tensor:
unique = np.unique(face_chart[face_chart >= 0])
if unique.size == 0:
return torch.from_numpy(face_chart).to(device)
remap = np.full(int(unique.max()) + 1, -1, dtype=np.int64)
remap[unique] = np.arange(unique.size)
out = face_chart.copy()
mask = out >= 0
out[mask] = remap[out[mask]]
return torch.from_numpy(out).to(device)
def segment_charts( def segment_charts(
@ -166,162 +89,108 @@ def segment_charts(
if F == 0: if F == 0:
return torch.zeros(0, dtype=torch.long, device=device) return torch.zeros(0, dtype=torch.long, device=device)
face_normal = mesh.face_normal.detach().cpu().numpy().astype(np.float32) fn = mesh.face_normal.detach().to(torch.float32)
face_area = mesh.face_area.detach().cpu().numpy().astype(np.float32) fa = mesh.face_area.detach().to(torch.float32)
face_centroid = mesh.face_centroid.detach().cpu().numpy().astype(np.float32) fc = mesh.face_centroid.detach().to(torch.float32)
face_face = mesh.face_face.detach().cpu().numpy() ff = mesh.face_face.detach().long()
fel = face_edge_lengths(mesh.vertices, mesh.faces).detach().to(torch.float32)
nd_cutoff = NORMAL_DEVIATION_HARD_CUTOFF
face_chart = np.full(F, -1, dtype=np.int64) # one seed per connected component (first face of each)
nd_cutoff = np.float32(NORMAL_DEVIATION_HARD_CUTOFF) comp = mesh.component.detach().long().to(device)
nd_threshold = np.float32(min(max_cost / max(w_normal_deviation, 1e-6), ncomp = int(comp.max()) + 1 if comp.numel() else 0
NORMAL_DEVIATION_HARD_CUTOFF * 0.99)) if ncomp:
seeds = torch.full((ncomp,), F, dtype=torch.long, device=device)
component = (mesh.component.detach().cpu().numpy() seeds.scatter_reduce_(0, comp, torch.arange(F, device=device), reduce="amin")
if hasattr(mesh.component, "detach") else np.asarray(mesh.component))
if component.size:
_, first_idx = np.unique(component, return_index=True)
initial_seeds = first_idx.astype(np.int64)
else: else:
initial_seeds = np.empty(0, dtype=np.int64) seeds = torch.zeros(1, dtype=torch.long, device=device)
K = seeds.shape[0]
seed_faces: List[int] = [int(s) for s in initial_seeds.tolist()] max_total_charts = max(F, 8000)
if not seed_faces: cap = K + F + 1 # every re-seed assigns a face, so K < K0 + F
seed_faces = [0] face_chart = torch.full((F,), -1, dtype=torch.long, device=device)
basis = torch.zeros(cap, 3, dtype=torch.float32, device=device)
nsum = torch.zeros(cap, 3, dtype=torch.float32, device=device)
area = torch.zeros(cap, dtype=torch.float32, device=device)
perim = torch.zeros(cap, dtype=torch.float32, device=device)
face_chart[seeds] = torch.arange(K, device=device)
basis[:K] = fn[seeds]
nsum[:K] = fn[seeds] * fa[seeds, None]
area[:K] = fa[seeds]
perim[:K] = fel[seeds].sum(1)
frontier = torch.zeros(F, dtype=torch.bool, device=device)
seed_nb = ff[seeds]
seed_nb = seed_nb[seed_nb >= 0]
frontier[seed_nb] = True
frontier &= face_chart < 0
K = len(seed_faces) min_d2 = torch.full((F,), float("inf"), dtype=torch.float32, device=device)
chart_basis = np.zeros((K, 3), dtype=np.float32) for i in range(0, K, 32): # chunked: (F, <=32, 3) stays small
chart_normal_sum = np.zeros((K, 3), dtype=np.float32) d2 = ((fc[:, None, :] - fc[seeds[i:i + 32]][None, :, :]) ** 2).sum(-1)
chart_area = np.zeros(K, dtype=np.float32) min_d2 = torch.minimum(min_d2, d2.amin(1))
chart_perim = np.zeros(K, dtype=np.float32)
face_edge_len = (
face_edge_lengths(mesh.vertices, mesh.faces)
.detach().cpu().numpy()
)
for cid, sf in enumerate(seed_faces):
face_chart[sf] = cid
n = face_normal[sf]
a = face_area[sf]
chart_basis[cid] = n.astype(np.float32)
chart_normal_sum[cid] = (n * a).astype(np.float32)
chart_area[cid] = float(a)
chart_perim[cid] = float(face_edge_len[sf].sum())
if K == 0: # Multi-pass threshold schedule (low-cost first); tau cap 0.5 keeps cones ~30deg.
return _renumber(face_chart, device) tau_final = min(max_cost * 0.25, 0.5)
thresholds = [t for t in (0.05, 0.1, 0.25) if t < tau_final] + [tau_final]
max_inner = max(64, int(F ** 0.5) * 2)
outer_iter = 0
tq = tqdm(total=F, desc="unwrap: segment (adaptive)", unit="face", leave=False)
while True:
outer_iter += 1
if outer_iter > F + 16:
break
for tau in thresholds:
for _ in range(max_inner):
n_added = _grow_iter(face_chart, frontier, ff, fn, fa, fel, basis, nsum,
area, perim, K, nd_cutoff, tau, w_normal_deviation,
w_roundness, w_straightness)
if n_added == 0:
break
tq.update(n_added)
unassigned = face_chart < 0
if int(unassigned.sum()) == 0:
break
if K >= max_total_charts:
break
# re-seed at the unassigned face farthest from every existing seed
new_seed = int(torch.where(unassigned, min_d2,
min_d2.new_full((), float("-inf"))).argmax())
face_chart[new_seed] = K
basis[K] = fn[new_seed]
nsum[K] = fn[new_seed] * fa[new_seed]
area[K] = fa[new_seed]
perim[K] = fel[new_seed].sum()
K += 1
min_d2 = torch.minimum(min_d2, ((fc - fc[new_seed]) ** 2).sum(-1))
tq.update(1)
frontier[new_seed] = False
ns_nb = ff[new_seed]
ns_nb = ns_nb[ns_nb >= 0]
frontier[ns_nb[face_chart[ns_nb] < 0]] = True
min_dist_to_seed = np.full(F, np.inf, dtype=np.float32) tq.close()
for sf in seed_faces:
d = ((face_centroid - face_centroid[sf]) ** 2).sum(axis=-1)
min_dist_to_seed = np.minimum(min_dist_to_seed, d)
if _HAVE_NUMBA:
# Multi-pass threshold schedule (low-cost first); tau cap 0.5 keeps cones ~30deg.
tau_final = min(max_cost * 0.25, 0.5)
thresholds = [t for t in (0.05, 0.1, 0.25) if t < tau_final] + [tau_final]
max_inner = max(64, int(np.sqrt(F)) * 2)
max_total_charts = max(F, 8000)
outer_iter = 0
while True:
outer_iter += 1
if outer_iter > F + 16:
break
for tau in thresholds:
for _ in range(max_inner):
n_added = _cost_grow_iter_jit(
face_chart, face_face, face_normal, face_area, face_edge_len,
chart_basis, chart_normal_sum, chart_area, chart_perim,
nd_cutoff, np.float32(tau),
np.float32(w_normal_deviation),
np.float32(w_roundness),
np.float32(w_straightness),
)
if n_added == 0:
break
if (face_chart == -1).sum() == 0:
break
if chart_basis.shape[0] >= max_total_charts:
break
unassigned_mask = face_chart == -1
cand = np.where(unassigned_mask, min_dist_to_seed, np.float32(-np.inf))
new_seed = int(np.argmax(cand))
n = face_normal[new_seed]
a = face_area[new_seed]
chart_basis = np.vstack([chart_basis, n[None, :].astype(np.float32)])
chart_normal_sum = np.vstack(
[chart_normal_sum, (n * a)[None, :].astype(np.float32)]
)
chart_area = np.concatenate([chart_area, np.array([a], dtype=np.float32)])
chart_perim = np.concatenate(
[chart_perim, np.array([face_edge_len[new_seed].sum()], dtype=np.float32)]
)
face_chart[new_seed] = chart_basis.shape[0] - 1
new_d = ((face_centroid - face_centroid[new_seed]) ** 2).sum(axis=-1)
min_dist_to_seed = np.minimum(min_dist_to_seed, new_d)
else:
# Numpy fallback: nd-only adaptive grow.
for _ in range(max(64, int(np.sqrt(F)) + 32)):
unassigned = face_chart == -1
if not unassigned.any():
break
u_idx = np.nonzero(unassigned)[0]
nbs = face_face[u_idx]
nbs_safe = np.where(nbs >= 0, nbs, 0)
nb_charts = np.where(nbs >= 0, face_chart[nbs_safe], -1)
valid = (nb_charts >= 0)
if not valid.any():
break
nb_charts_safe = np.where(valid, nb_charts, 0)
nb_basis = chart_basis[nb_charts_safe]
d = (face_normal[u_idx][:, None, :] * nb_basis).sum(axis=-1)
nd = np.where(valid, np.float32(1.0) - d, np.inf).clip(max=1.0)
nd = np.where(nd >= nd_cutoff, np.inf, nd)
best_e = np.argmin(nd, axis=1)
best_cost = nd[np.arange(u_idx.size), best_e]
best_c = nb_charts_safe[np.arange(u_idx.size), best_e]
accept = (best_cost <= nd_threshold) & np.isfinite(best_cost)
if not accept.any():
break
pick_u = u_idx[accept]
pick_c = best_c[accept]
face_chart[pick_u] = pick_c
for f, c in zip(pick_u, pick_c):
chart_normal_sum[c] += face_normal[f] * face_area[f]
chart_area[c] += face_area[f]
# Orphan cleanup: leftover faces join their best-matching neighbor's chart. # Orphan cleanup: leftover faces join their best-matching neighbor's chart.
if (face_chart == -1).any() and chart_basis.shape[0] > 0: while True:
while True: orphans = (face_chart < 0).nonzero(as_tuple=True)[0]
orphans = np.nonzero(face_chart == -1)[0] if orphans.numel() == 0:
if orphans.size == 0: break
break nb = ff[orphans]
nbs = face_face[orphans] nbc = torch.where(nb >= 0, face_chart[nb.clamp_min(0)], nb.new_full((), -1))
nbs_safe = np.where(nbs >= 0, nbs, 0) valid = nbc >= 0
nb_charts = np.where(nbs >= 0, face_chart[nbs_safe], -1) assignable = valid.any(1)
valid = (nb_charts >= 0) if not bool(assignable.any()):
if not valid.any(): break
break d = (fn[orphans][:, None, :] * basis[nbc.clamp_min(0)]).sum(-1)
nb_charts_safe = np.where(valid, nb_charts, 0) ndv = torch.where(valid, 1.0 - d, d.new_full((), float("inf")))
nb_basis = chart_basis[nb_charts_safe] best_c = nbc.gather(1, ndv.argmin(1, keepdim=True)).squeeze(1)
d = (face_normal[orphans][:, None, :] * nb_basis).sum(axis=-1) face_chart[orphans[assignable]] = best_c[assignable]
nd = np.where(valid, np.float32(1.0) - d, np.inf) leftover = (face_chart < 0).nonzero(as_tuple=True)[0]
best_e = np.argmin(nd, axis=1) if leftover.numel(): # isolated faces become singleton charts
best_c = nb_charts_safe[np.arange(orphans.size), best_e] face_chart[leftover] = K + torch.arange(leftover.numel(), device=device)
assignable = valid.any(axis=1)
if not assignable.any():
break
assign_idx = orphans[assignable]
assign_c = best_c[assignable]
face_chart[assign_idx] = assign_c
if (face_chart == -1).any():
new_singletons = np.nonzero(face_chart == -1)[0]
for f in new_singletons:
face_chart[int(f)] = chart_basis.shape[0]
chart_basis = np.concatenate(
[chart_basis, face_normal[int(f)].astype(np.float32)[None, :]],
axis=0,
)
return _renumber(face_chart, device) _, inverse = torch.unique(face_chart, sorted=True, return_inverse=True)
return inverse
# Parallel edge-collapse (PEC) chart clustering (GPU) # Parallel edge-collapse (PEC) chart clustering (GPU)
@ -400,12 +269,71 @@ def _build_chart_edges(
return chart_pairs, reduced_el return chart_pairs, reduced_el
def _merge_small_charts(
chart_id: Tensor, face_normal: Tensor, face_area: Tensor,
face_face: Tensor, face_edge_len: Tensor,
min_faces: int, cost_cap: float,
) -> Tensor:
"""Absorb charts under min_faces faces into their lowest-cone-cost neighbor (capped at cost_cap)."""
if chart_id.numel() == 0:
return chart_id
device = chart_id.device
for _ in range(16):
N = int(chart_id.max().item()) + 1
sizes = torch.bincount(chart_id, minlength=N)
# recompute cones from scratch: area-weighted mean axis, max deviation as half-angle
axis = torch.zeros(N, 3, dtype=torch.float32, device=device)
axis.index_add_(0, chart_id, face_normal * face_area[:, None])
axis = axis / axis.norm(dim=1, keepdim=True).clamp_min(1e-12)
dev = torch.acos((face_normal * axis[chart_id]).sum(1).clamp(-1.0, 1.0))
half = torch.zeros(N, dtype=torch.float32, device=device)
half.scatter_reduce_(0, chart_id, dev, reduce="amax")
edges, _ = _build_chart_edges(face_face, chart_id, face_edge_len)
if edges.shape[0] == 0:
break
a, b = edges[:, 0], edges[:, 1]
_, new_half, _ = _combine_normal_cones(axis[a], half[a], axis[b], half[b])
ok = new_half <= cost_cap
E = edges.shape[0]
key = (torch.clamp(new_half * 1e6, max=2e9).to(torch.int64) << 32) \
| torch.arange(E, dtype=torch.long, device=device)
best = torch.full((N,), 1 << 62, dtype=torch.long, device=device)
va = (sizes[a] < min_faces) & ok
vb = (sizes[b] < min_faces) & ok
best.scatter_reduce_(0, a[va], key[va], reduce="amin")
best.scatter_reduce_(0, b[vb], key[vb], reduce="amin")
src = (best < (1 << 62)).nonzero(as_tuple=True)[0]
if src.numel() == 0:
break
eid = best[src] & 0xFFFFFFFF
ea, eb = a[eid], b[eid]
tgt = torch.where(ea == src, eb, ea)
# cycle break: keep src->tgt only if tgt merges nowhere itself or src > tgt;
# the kept graph is then a DAG, so the pointer-doubling below terminates
prop = torch.arange(N, dtype=torch.long, device=device)
prop[src] = tgt
keepm = (prop[tgt] == tgt) | (src > tgt)
remap = torch.arange(N, dtype=torch.long, device=device)
remap[src[keepm]] = tgt[keepm]
for _ in range(32):
nr = remap[remap]
if torch.equal(nr, remap):
break
remap = nr
chart_id = remap[chart_id]
_, chart_id = torch.unique(chart_id, return_inverse=True)
return chart_id
def cluster_charts_pec( def cluster_charts_pec(
mesh: MeshData, mesh: MeshData,
max_cost: float = 0.7, max_cost: float = 0.7,
max_iters: int = 1024, max_iters: int = 1024,
min_faces: int = 8,
) -> Tensor: ) -> Tensor:
"""Parallel edge-collapse clustering; returns face_chart [F]. max_cost is the per-merge cutoff (~0.7 rad ~ 40deg).""" """Parallel edge-collapse clustering; returns face_chart [F]. max_cost is the per-merge
cutoff (~0.7 rad ~ 40deg); charts under min_faces are then absorbed at a relaxed 2x cutoff."""
device = mesh.faces.device device = mesh.faces.device
F = mesh.faces.shape[0] F = mesh.faces.shape[0]
faces = mesh.faces.to(torch.long) faces = mesh.faces.to(torch.long)
@ -471,5 +399,8 @@ def cluster_charts_pec(
remap[win_b] = win_a remap[win_b] = win_a
chart_id = remap[chart_id] chart_id = remap[chart_id]
if min_faces > 1:
chart_id = _merge_small_charts(chart_id, face_normal, mesh.face_area.to(torch.float32),
face_face, face_edge_len, min_faces, 2.0 * max_cost)
_, inverse = torch.unique(chart_id, sorted=True, return_inverse=True) _, inverse = torch.unique(chart_id, sorted=True, return_inverse=True)
return inverse return inverse

166
comfy_extras/nodes_mesh_postprocess.py
View File

@ -15,6 +15,7 @@ from comfy_extras.mesh3d.uv_unwrap import segment as _uv_seg
from comfy_extras.mesh3d.uv_unwrap import parameterize as _uv_param from comfy_extras.mesh3d.uv_unwrap import parameterize as _uv_param
from comfy_extras.mesh3d.uv_unwrap import pack as _uv_pack from comfy_extras.mesh3d.uv_unwrap import pack as _uv_pack
import logging import logging
import time
from tqdm import tqdm from tqdm import tqdm
from scipy.sparse import csr_matrix from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import connected_components from scipy.sparse.csgraph import connected_components
@ -2454,6 +2455,7 @@ def _uv_weld_vertices(v, f, weld_distance):
def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance): def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance):
"""UV-unwrap a single mesh; returns (vmapping, indices, uvs); vmapping maps each output """UV-unwrap a single mesh; returns (vmapping, indices, uvs); vmapping maps each output
vertex to an input vertex (seam verts duplicated).""" vertex to an input vertex (seam verts duplicated)."""
t_start = time.perf_counter()
v_in = positions.to(torch.float32) v_in = positions.to(torch.float32)
f_in = indices.to(torch.long).reshape(-1, 3) f_in = indices.to(torch.long).reshape(-1, 3)
v_in, f_in, welded_to_orig = _uv_weld_vertices(v_in, f_in, weld_distance) v_in, f_in, welded_to_orig = _uv_weld_vertices(v_in, f_in, weld_distance)
@ -2479,72 +2481,138 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance
n_charts = int(face_chart.max().item()) + 1 if face_chart.numel() else 0 n_charts = int(face_chart.max().item()) + 1 if face_chart.numel() else 0
areas_cpu = _uv_mesh.chart_3d_areas(mesh.face_area, face_chart, n_charts).detach().cpu() areas_cpu = _uv_mesh.chart_3d_areas(mesh.face_area, face_chart, n_charts).detach().cpu()
# per-chart loop on CPU/numpy to avoid per-chart GPU sync if n_charts == 0:
return (np.empty(0, dtype=np.int64), np.zeros((0, 3), dtype=np.int64),
np.empty((0, 2), dtype=np.float32))
# vectorized chart extraction: one global sort/unique replaces per-chart unique/searchsorted
face_chart_np = face_chart.cpu().numpy() face_chart_np = face_chart.cpu().numpy()
faces_np = mesh.faces.cpu().numpy() faces_np = mesh.faces.cpu().numpy()
vertices_np = mesh.vertices.cpu().numpy() vertices_np = mesh.vertices.cpu().numpy()
face_face_np = mesh.face_face.cpu().numpy() face_face_np = mesh.face_face.cpu().numpy()
sorted_face_idx_np = np.argsort(face_chart_np, kind="stable") order = np.argsort(face_chart_np, kind="stable")
chart_counts_np = np.bincount(face_chart_np, minlength=n_charts) chart_counts_np = np.bincount(face_chart_np, minlength=n_charts)
chart_offsets_np = np.empty(n_charts + 1, dtype=np.int64) chart_offsets_np = np.zeros(n_charts + 1, dtype=np.int64)
chart_offsets_np[0] = 0
np.cumsum(chart_counts_np, out=chart_offsets_np[1:]) np.cumsum(chart_counts_np, out=chart_offsets_np[1:])
faces_sorted = faces_np[order]
chart_sorted = face_chart_np[order]
n_verts_in = max(vertices_np.shape[0], 1)
chart_of_slot = np.repeat(chart_sorted, 3)
uniq_keys, local_flat = np.unique(chart_of_slot * n_verts_in + faces_sorted.reshape(-1),
return_inverse=True)
used_verts_all = uniq_keys % n_verts_in # per-chart sorted unique verts, concatenated
vert_counts = np.bincount(uniq_keys // n_verts_in, minlength=n_charts)
vert_offsets = np.zeros(n_charts + 1, dtype=np.int64)
np.cumsum(vert_counts, out=vert_offsets[1:])
local_faces_all = (local_flat - vert_offsets[chart_of_slot]).reshape(-1, 3)
pos_in_chart = np.empty(order.size, dtype=np.int64)
pos_in_chart[order] = np.arange(order.size) - chart_offsets_np[chart_sorted]
ff_sorted = face_face_np[order]
ff_safe = np.maximum(ff_sorted, 0)
keep = (ff_sorted >= 0) & (face_chart_np[ff_safe] == chart_sorted[:, None])
local_ff_all = np.where(keep, pos_in_chart[ff_safe], -1)
all_chart_uvs, all_chart_3d_areas, all_chart_uv_areas, all_chart_faces = [], [], [], [] # progress: n_charts units for parameterize + 2*n_charts for pack (prepare + place)
chart_records = [] pbar = comfy.utils.ProgressBar(3 * n_charts)
# parameterize (batched): ortho-project every chart at once, batched stretch metrics
# decide acceptance, rejected charts solve ABF/LSCM in dense per-size-bucket batches
chart_of_vert = (uniq_keys // n_verts_in).astype(np.int64)
verts_concat = vertices_np[used_verts_all].astype(np.float64)
gl_faces = local_faces_all + vert_offsets[chart_sorted][:, None]
face_pos = pos_in_chart[order] # row of each (sorted) face in its chart
uv0 = _uv_param.ortho_project_concat(verts_concat, chart_of_vert, n_charts)
rms, mx, n_flip, n_zero = _uv_param.stretch_metrics_concat(
verts_concat, uv0, gl_faces, chart_sorted, n_charts)
valid_chart = (vert_counts >= 3) & (chart_counts_np > 0)
auto = valid_chart & (chart_counts_np <= 5) # tiny charts always keep ortho
flip_ok = (n_flip == 0) | (n_flip == chart_counts_np)
cand = valid_chart & ~auto & flip_ok & (n_zero == 0) & (rms <= 1.5) & (mx <= 2.0)
pbar.update(int(auto.sum()))
ortho_ok = auto.copy()
cand_ids = np.nonzero(cand)[0]
for c in tqdm(cand_ids, desc="unwrap: ortho checks", unit="chart", leave=False):
f0, f1 = chart_offsets_np[c], chart_offsets_np[c + 1]
v0, v1 = vert_offsets[c], vert_offsets[c + 1]
if not _uv_param._uv_boundary_self_intersects(
uv0[v0:v1], local_faces_all[f0:f1], local_ff_all[f0:f1]):
ortho_ok[c] = True
pbar.update(1)
lscm_mask = valid_chart & ~ortho_ok
batchable = vert_counts <= _uv_param.LSCM_BATCH_MAX_VERTS
lscm_ids = np.nonzero(lscm_mask & batchable)[0]
big_ids = np.nonzero(lscm_mask & ~batchable)[0]
lscm_uv = _uv_param.lscm_charts_batch(
verts_concat, uv0, gl_faces, face_pos, chart_sorted, chart_of_vert,
vert_offsets, lscm_ids, n_charts,
device=comfy.model_management.get_torch_device())
pbar.update(int(lscm_ids.size))
uvs_np_list: list = [None] * n_charts
uv0_f32 = uv0.astype(np.float32)
for c in tqdm(big_ids, desc="unwrap: LSCM (large charts)", unit="chart", leave=False):
f0, f1 = chart_offsets_np[c], chart_offsets_np[c + 1]
v0, v1 = vert_offsets[c], vert_offsets[c + 1]
uvs_t = _uv_param.lscm_chart(
torch.from_numpy(verts_concat[v0:v1]),
torch.from_numpy(local_faces_all[f0:f1]),
torch.from_numpy(local_ff_all[f0:f1]), pin_positions=uv0[v0:v1])
lscm_uv[int(c)] = uvs_t.detach().cpu().numpy().astype(np.float32)
pbar.update(1)
for c in range(n_charts): for c in range(n_charts):
gfi_np = sorted_face_idx_np[chart_offsets_np[c]:chart_offsets_np[c + 1]] v0, v1 = vert_offsets[c], vert_offsets[c + 1]
chart_faces_global = faces_np[gfi_np] if ortho_ok[c]:
used_verts_np = np.unique(chart_faces_global) uvs_np_list[c] = uv0_f32[v0:v1]
local_faces_np = np.searchsorted(used_verts_np, chart_faces_global) continue
local_verts_np = vertices_np[used_verts_np] u = lscm_uv.get(int(c))
ff_global = face_face_np[gfi_np] if u is not None and np.all(np.isfinite(u)) and u.size:
ff_safe = np.maximum(ff_global, 0) # collapsed UV island (aspect > 100:1) blows up packing scale; keep ortho instead
nb_chart = np.where(ff_global >= 0, face_chart_np[ff_safe], -1) bbox = u.max(axis=0) - u.min(axis=0)
keep = (ff_global >= 0) & (nb_chart == c) if max(float(bbox.max()), 1e-12) / max(float(bbox.min()), 1e-12) <= 100.0:
local_neighbor = np.searchsorted(gfi_np, ff_safe) uvs_np_list[c] = u
local_ff_np = np.where(keep, local_neighbor, -1) continue
uvs_np_list[c] = (uv0_f32[v0:v1] if valid_chart[c]
else np.zeros((v1 - v0, 2), dtype=np.float32))
lf = torch.from_numpy(local_faces_np) # per-chart UV areas in one pass over all faces
uvs = _uv_param.parametrize_chart( uvs_all_np = np.concatenate(uvs_np_list)
torch.from_numpy(local_verts_np), lf, torch.from_numpy(local_ff_np)) ua, ub, uc = uvs_all_np[gl_faces[:, 0]], uvs_all_np[gl_faces[:, 1]], uvs_all_np[gl_faces[:, 2]]
ua, ub, uc = uvs[lf[:, 0]], uvs[lf[:, 1]], uvs[lf[:, 2]] tri_uv_area = 0.5 * np.abs(
uv_area_sum = float(0.5 * ( (ub[:, 0] - ua[:, 0]) * (uc[:, 1] - ua[:, 1])
(ub[:, 0] - ua[:, 0]) * (uc[:, 1] - ua[:, 1]) - (uc[:, 0] - ua[:, 0]) * (ub[:, 1] - ua[:, 1]))
- (uc[:, 0] - ua[:, 0]) * (ub[:, 1] - ua[:, 1])).abs().sum().item()) uv_area_np = np.bincount(chart_sorted, weights=tri_uv_area.astype(np.float64),
chart_records.append({"local_faces": lf, "vmap": torch.from_numpy(used_verts_np), minlength=n_charts)
"global_face_idx": torch.from_numpy(gfi_np)})
all_chart_uvs.append(uvs) areas_3d_np = areas_cpu.numpy().astype(np.float64)
all_chart_3d_areas.append(float(areas_cpu[c].item()))
all_chart_uv_areas.append(uv_area_sum)
all_chart_faces.append(lf)
# auto-tune texel density toward `resolution` (~0.62 pack fill) # auto-tune texel density toward `resolution` (~0.62 pack fill)
total_3d_area = sum(all_chart_3d_areas) or 1.0 total_3d_area = float(areas_3d_np.sum()) or 1.0
target_dim = float(resolution) if resolution > 0 else 1024.0 target_dim = float(resolution) if resolution > 0 else 1024.0
tex_per_unit = math.sqrt((target_dim * target_dim) * 0.62 / total_3d_area) tex_per_unit = math.sqrt((target_dim * target_dim) * 0.62 / total_3d_area)
placements, atlas_w, atlas_h = _uv_pack.pack_bitmap( with tqdm(total=2 * n_charts, desc="unwrap: pack", unit="chart", leave=False) as tq_pack:
all_chart_uvs, all_chart_3d_areas, all_chart_uv_areas, all_chart_faces, def _pack_progress(done, total):
texels_per_unit=tex_per_unit, padding_texels=padding) tq_pack.update(done - tq_pack.n)
placed = _uv_pack.apply_placements(all_chart_uvs, placements, atlas_w, atlas_h) pbar.update_absolute(n_charts + done, 3 * n_charts)
p_x, p_y, p_sw, p_th, p_sc, p_chh, atlas_w, atlas_h = _uv_pack.pack_bitmap_concat(
uvs_all_np, vert_offsets, local_faces_all, chart_offsets_np,
areas_3d_np, uv_area_np,
texels_per_unit=tex_per_unit, padding_texels=padding,
progress_callback=_pack_progress)
pbar.update_absolute(3 * n_charts, 3 * n_charts)
# assembly: output verts are the per-chart used-vert lists concatenated in chart order,
# so vert_offsets doubles as the output vertex cursor
n_in_faces = mesh.faces.shape[0] n_in_faces = mesh.faces.shape[0]
out_indices = np.zeros((n_in_faces, 3), dtype=np.int64) out_indices = np.zeros((n_in_faces, 3), dtype=np.int64)
out_uvs_list, out_vmap_list, v_cursor = [], [], 0 out_indices[order] = gl_faces
for c, rec in enumerate(chart_records): vmapping_out = used_verts_all if welded_to_orig is None else welded_to_orig[used_verts_all]
vmap_np = rec["vmap"].cpu().numpy() uvs_out = _uv_pack.apply_placements_concat(
local_faces_np = rec["local_faces"].cpu().numpy() uvs_all_np, vert_offsets, p_x, p_y, p_sw, p_th, p_sc, p_chh, atlas_w, atlas_h)
global_face_idx = rec["global_face_idx"].cpu().numpy() logging.info(f"[uv_unwrap] {mesh.faces.shape[0]} faces -> {n_charts} charts, "
out_uvs_list.append(placed[c].cpu().numpy()) f"atlas {atlas_w}x{atlas_h}, {time.perf_counter() - t_start:.1f}s")
if welded_to_orig is not None:
vmap_np = welded_to_orig[vmap_np]
out_vmap_list.append(vmap_np)
out_indices[global_face_idx] = local_faces_np + v_cursor
v_cursor += vmap_np.shape[0]
vmapping_out = np.concatenate(out_vmap_list) if out_vmap_list else np.empty(0, dtype=np.int64)
uvs_out = np.concatenate(out_uvs_list) if out_uvs_list else np.empty((0, 2), dtype=np.float32)
return vmapping_out, out_indices, uvs_out return vmapping_out, out_indices, uvs_out