From 5067ac461ecb328db438579355e20008ec51b0a0 Mon Sep 17 00:00:00 2001 From: kijai Date: Fri, 3 Jul 2026 20:23:26 +0300 Subject: [PATCH] Refactor/optimize UV unwrap --- comfy_extras/mesh3d/uv_unwrap/pack.py | 1201 +++++++++-------- comfy_extras/mesh3d/uv_unwrap/parameterize.py | 316 ++++- comfy_extras/mesh3d/uv_unwrap/segment.py | 499 +++---- comfy_extras/nodes_mesh_postprocess.py | 166 ++- 4 files changed, 1196 insertions(+), 986 deletions(-) diff --git a/comfy_extras/mesh3d/uv_unwrap/pack.py b/comfy_extras/mesh3d/uv_unwrap/pack.py index 5b49b613f..219f5e7e0 100644 --- a/comfy_extras/mesh3d/uv_unwrap/pack.py +++ b/comfy_extras/mesh3d/uv_unwrap/pack.py @@ -3,24 +3,32 @@ from __future__ import annotations import math from dataclasses import dataclass -from typing import List, Tuple +from typing import Tuple import numpy as np import torch from torch import Tensor +from torch.nn.functional import max_pool1d import comfy.model_management +# Numba is optional, but ~5x faster than torch on these operations, potential TODO: comfy-kitchen cuda/triton kernels as even faster alternative try: - from numba import njit as _njit + from numba import njit as _njit, prange as _prange, get_num_threads as _nb_threads _HAVE_NUMBA_PACK = True except ImportError: _HAVE_NUMBA_PACK = False + _prange = range + def _nb_threads(): return 1 def _njit(*args, **kwargs): def deco(fn): return fn return deco if not args else args[0] +# Cap on deterministic sweep density: tiny charts on a large atlas would otherwise enumerate every texel column. +_SWEEP_CAP = 1024 + + @dataclass class ChartPlacement: chart_id: int @@ -31,27 +39,50 @@ class ChartPlacement: chart_h: float = 0.0 # unswapped bitmap height in texels (rotation pivot) -@_njit(cache=True, boundscheck=False) -def _best_rotation_jit(uvs_np: np.ndarray, n_angles: int) -> float: - V = uvs_np.shape[0] - best_area = 1e30 - best_theta = 0.0 - if V == 0: - return 0.0 +@_njit(cache=True, boundscheck=False, parallel=True) +def _prepare_dims_jit(uvs, uv_off, a3, auv, tpu, padding, theta, scale, bw, bh, rot_uv): + """Pass 1: per-chart best rotation, texel scale, rotated/scaled UVs, padded bitmap dims.""" + n = uv_off.shape[0] - 1 half_pi = math.pi * 0.5 - for k in range(n_angles): - theta = half_pi * k / n_angles - c = math.cos(theta) - s = math.sin(theta) + for c in _prange(n): + v0, v1 = uv_off[c], uv_off[c + 1] + best_area = 1e30 + best_t = 0.0 + for k in range(36): + th = half_pi * k / 36.0 + co = math.cos(th) + si = math.sin(th) + xmin = 1e30 + xmax = -1e30 + ymin = 1e30 + ymax = -1e30 + for i in range(v0, v1): + xr = uvs[i, 0] * co - uvs[i, 1] * si + yr = uvs[i, 0] * si + uvs[i, 1] * co + if xr < xmin: + xmin = xr + if xr > xmax: + xmax = xr + if yr < ymin: + ymin = yr + if yr > ymax: + ymax = yr + area = (xmax - xmin) * (ymax - ymin) + if area < best_area: + best_area = area + best_t = th + theta[c] = best_t + co = math.cos(best_t) + si = math.sin(best_t) xmin = 1e30 xmax = -1e30 ymin = 1e30 ymax = -1e30 - for i in range(V): - ux = uvs_np[i, 0] - uy = uvs_np[i, 1] - xr = ux * c - uy * s - yr = ux * s + uy * c + for i in range(v0, v1): + xr = uvs[i, 0] * co - uvs[i, 1] * si + yr = uvs[i, 0] * si + uvs[i, 1] * co + rot_uv[i, 0] = xr + rot_uv[i, 1] = yr if xr < xmin: xmin = xr if xr > xmax: @@ -60,315 +91,334 @@ def _best_rotation_jit(uvs_np: np.ndarray, n_angles: int) -> float: ymin = yr if yr > ymax: ymax = yr - area = (xmax - xmin) * (ymax - ymin) - if area < best_area: - best_area = area - best_theta = theta - return best_theta + if v1 == v0: + xmin = 0.0 + xmax = 0.0 + ymin = 0.0 + ymax = 0.0 + s = math.sqrt(max(a3[c], 1e-12) / max(auv[c], 1e-12)) * tpu + nominal = math.sqrt(max(a3[c], 1e-12)) * tpu + max_bbox = max(8.0, 4.0 * nominal) + bbox_max = max(max(xmax - xmin, ymax - ymin), 1e-12) + if s * bbox_max > max_bbox: + s = max_bbox / bbox_max + scale[c] = s + wmax = 0.0 + hmax = 0.0 + for i in range(v0, v1): + rot_uv[i, 0] = (rot_uv[i, 0] - xmin) * s + rot_uv[i, 1] = (rot_uv[i, 1] - ymin) * s + if rot_uv[i, 0] > wmax: + wmax = rot_uv[i, 0] + if rot_uv[i, 1] > hmax: + hmax = rot_uv[i, 1] + bw[c] = int(math.ceil(wmax)) + padding + 1 + bh[c] = int(math.ceil(hmax)) + padding + 1 -def _best_rotation(uvs_np: np.ndarray, n_angles: int = 36) -> float: - return float(_best_rotation_jit(uvs_np.astype(np.float64), n_angles)) - - -def _rotate_xy(uv: np.ndarray, theta: float) -> np.ndarray: - if theta == 0.0: - return uv - c = math.cos(theta) - s = math.sin(theta) - return np.stack([uv[:, 0] * c - uv[:, 1] * s, uv[:, 0] * s + uv[:, 1] * c], axis=1) - - -@_njit(cache=True, boundscheck=False) -def _rasterize_chart_jit( - uvs_tex: np.ndarray, faces: np.ndarray, w: int, h: int -) -> np.ndarray: - """JIT-rasterize triangles into an (h, w) bool bitmap via barycentric test.""" - bm = np.zeros((h, w), dtype=np.bool_) - F = faces.shape[0] +@_njit(cache=True, boundscheck=False, parallel=True) +def _raster_all_jit(rot_uv, uv_off, faces, f_off, bw, bh, boff, buf, padding, + tw, th_out, perim): + """Pass 2: rasterize + dilate each chart into the flat buffer; records trimmed dims + (origin kept) and the perimeter used for placement ordering.""" + n = uv_off.shape[0] - 1 eps = 1e-7 - for fi in range(F): - i0 = faces[fi, 0] - i1 = faces[fi, 1] - i2 = faces[fi, 2] - x0 = uvs_tex[i0, 0] - y0 = uvs_tex[i0, 1] - x1 = uvs_tex[i1, 0] - y1 = uvs_tex[i1, 1] - x2 = uvs_tex[i2, 0] - y2 = uvs_tex[i2, 1] - xmin_f = x0 - if x1 < xmin_f: - xmin_f = x1 - if x2 < xmin_f: - xmin_f = x2 - xmax_f = x0 - if x1 > xmax_f: - xmax_f = x1 - if x2 > xmax_f: - xmax_f = x2 - ymin_f = y0 - if y1 < ymin_f: - ymin_f = y1 - if y2 < ymin_f: - ymin_f = y2 - ymax_f = y0 - if y1 > ymax_f: - ymax_f = y1 - if y2 > ymax_f: - ymax_f = y2 - xmin = int(math.floor(xmin_f)) - if xmin < 0: - xmin = 0 - xmax = int(math.ceil(xmax_f)) - if xmax > w - 1: - xmax = w - 1 - ymin = int(math.floor(ymin_f)) - if ymin < 0: - ymin = 0 - ymax = int(math.ceil(ymax_f)) - if ymax > h - 1: - ymax = h - 1 - if xmax < xmin or ymax < ymin: - continue - denom = (y1 - y2) * (x0 - x2) + (x2 - x1) * (y0 - y2) - if abs(denom) < 1e-20: - continue - inv_denom = 1.0 / denom - for py in range(ymin, ymax + 1): - yc = py + 0.5 - for px in range(xmin, xmax + 1): - xc = px + 0.5 - a = ((y1 - y2) * (xc - x2) + (x2 - x1) * (yc - y2)) * inv_denom - b = ((y2 - y0) * (xc - x2) + (x0 - x2) * (yc - y2)) * inv_denom - c = 1.0 - a - b - if a >= -eps and b >= -eps and c >= -eps: - bm[py, px] = True - return bm - - -def _rasterize_chart( - uvs_tex: np.ndarray, faces: np.ndarray, w: int, h: int, padding: int -) -> np.ndarray: - """Rasterize chart triangles into (h, w) bool bitmap, dilated by padding texels.""" - if faces.size == 0: - return np.zeros((h, w), dtype=bool) - bm = _rasterize_chart_jit( - uvs_tex.astype(np.float64), faces.astype(np.int64), int(w), int(h) - ) - if padding > 0: - bm = _dilate_bitmap(bm, padding) - return bm - - -def _dilate_bitmap(bm: np.ndarray, k: int) -> np.ndarray: - """k-step Manhattan max-filter dilation.""" - out = bm.copy() - for _ in range(k): - next_out = out.copy() - next_out[1:, :] |= out[:-1, :] - next_out[:-1, :] |= out[1:, :] - next_out[:, 1:] |= out[:, :-1] - next_out[:, :-1] |= out[:, 1:] - out = next_out - return out - - -@_njit(cache=True, boundscheck=False) -def _build_candidates_jit( - skyline: np.ndarray, - cur_w: int, cur_h: int, - bw0: int, bh0: int, bw1: int, bh1: int, - step: int, -) -> np.ndarray: - """Build per-chart (x, y, swap_flag) candidate positions (skyline-flush + edge-sweep, both orientations).""" - nx_skyline = (max(cur_w, 1) // step) + 2 - nx_edge = (max(cur_w, 1) // step) + 2 - ny_edge = (max(cur_h, 1) // step) + 2 - per_orient = nx_skyline + 2 * nx_edge + 2 * ny_edge - out = np.empty((per_orient * 2, 3), dtype=np.int64) - k = 0 - for swap_flag in range(2): - cw = bw0 if swap_flag == 0 else bw1 - x = 0 - while x <= cur_w: - y = 0 - x_end = x + cw - if x_end > skyline.shape[0]: - x_end = skyline.shape[0] - for xs in range(x, x_end): - if skyline[xs] > y: - y = int(skyline[xs]) - out[k, 0] = x - out[k, 1] = y - out[k, 2] = swap_flag - k += 1 - x += step - for y_fixed in (0, cur_h): - x = 0 - while x <= cur_w: - out[k, 0] = x - out[k, 1] = y_fixed - out[k, 2] = swap_flag - k += 1 - x += step - for x_fixed in (0, cur_w): - y = 0 - while y <= cur_h: - out[k, 0] = x_fixed - out[k, 1] = y - out[k, 2] = swap_flag - k += 1 - y += step - return out[:k] - - -@_njit(cache=True, boundscheck=False) -def _update_skyline_jit(skyline: np.ndarray, chart: np.ndarray, - x: int, y: int) -> None: - """Lift skyline[x+i] to y + topmost_True_row + 1 per chart column.""" - ch = chart.shape[0] - cw = chart.shape[1] - sw = skyline.shape[0] - for i in range(cw): - col_x = x + i - if col_x >= sw or col_x < 0: - continue - col_top = -1 - for j in range(ch - 1, -1, -1): - if chart[j, i]: - col_top = j - break - if col_top < 0: - continue - new_h = y + col_top + 1 - if new_h > skyline[col_x]: - skyline[col_x] = new_h - - -@_njit(cache=True, boundscheck=False) -def _best_placement_jit( - atlas: np.ndarray, - bitmap: np.ndarray, - bitmap_rot: np.ndarray, - candidates: np.ndarray, - cur_w: int, - cur_h: int, -): - """Pick lowest-score non-colliding candidate (score = max(new_w,new_h)^2 + new_w*new_h); out-of-atlas treated as free.""" - n = candidates.shape[0] - best_x = -1 - best_y = -1 - best_score = -1 - best_swap = 0 - bh0 = bitmap.shape[0] - bw0 = bitmap.shape[1] - bh1 = bitmap_rot.shape[0] - bw1 = bitmap_rot.shape[1] - ah = atlas.shape[0] - aw = atlas.shape[1] - for k in range(n): - x = candidates[k, 0] - y = candidates[k, 1] - swap = candidates[k, 2] - if swap == 0: - ch = bh0 - cw = bw0 - else: - ch = bh1 - cw = bw1 - if x < 0 or y < 0: - continue - nw = cur_w if cur_w > x + cw else x + cw - nh = cur_h if cur_h > y + ch else y + ch - ext = nw if nw > nh else nh - score = ext * ext + nw * nh - if best_score >= 0 and score >= best_score: - continue - ok = True - for j in range(ch): - yy = y + j - if yy >= ah: + for c in _prange(n): + f0, f1 = f_off[c], f_off[c + 1] + v0 = uv_off[c] + V = uv_off[c + 1] - v0 + w = bw[c] + h = bh[c] + o = boff[c] + for fi in range(f0, f1): + i0 = faces[fi, 0] + v0 + i1 = faces[fi, 1] + v0 + i2 = faces[fi, 2] + v0 + x0 = rot_uv[i0, 0] + y0 = rot_uv[i0, 1] + x1 = rot_uv[i1, 0] + y1 = rot_uv[i1, 1] + x2 = rot_uv[i2, 0] + y2 = rot_uv[i2, 1] + xmin_f = min(x0, min(x1, x2)) + xmax_f = max(x0, max(x1, x2)) + ymin_f = min(y0, min(y1, y2)) + ymax_f = max(y0, max(y1, y2)) + xmin = max(int(math.floor(xmin_f)), 0) + xmax = min(int(math.ceil(xmax_f)), w - 1) + ymin = max(int(math.floor(ymin_f)), 0) + ymax = min(int(math.ceil(ymax_f)), h - 1) + if xmax < xmin or ymax < ymin: continue - for i in range(cw): - bit = bitmap[j, i] if swap == 0 else bitmap_rot[j, i] - if not bit: + denom = (y1 - y2) * (x0 - x2) + (x2 - x1) * (y0 - y2) + if abs(denom) < 1e-20: + continue + inv_denom = 1.0 / denom + for py in range(ymin, ymax + 1): + yc = py + 0.5 + for px in range(xmin, xmax + 1): + xc = px + 0.5 + aa = ((y1 - y2) * (xc - x2) + (x2 - x1) * (yc - y2)) * inv_denom + bb = ((y2 - y0) * (xc - x2) + (x0 - x2) * (yc - y2)) * inv_denom + cc = 1.0 - aa - bb + if aa >= -eps and bb >= -eps and cc >= -eps: + buf[o + py * w + px] = True + # Manhattan dilation by `padding` steps (ping-pong on a scratch copy) + if padding > 0 and f1 > f0: + tmp = np.empty(h * w, dtype=np.bool_) + for _ in range(padding): + for j in range(h * w): + tmp[j] = buf[o + j] + for py in range(h): + for px in range(w): + if tmp[py * w + px]: + continue + hit = False + if py > 0 and tmp[(py - 1) * w + px]: + hit = True + elif py < h - 1 and tmp[(py + 1) * w + px]: + hit = True + elif px > 0 and tmp[py * w + px - 1]: + hit = True + elif px < w - 1 and tmp[py * w + px + 1]: + hit = True + if hit: + buf[o + py * w + px] = True + # trimmed dims (keep origin; 1x1 empty bitmap when nothing was rasterized) + rmax = -1 + cmax = -1 + for py in range(h): + for px in range(w): + if buf[o + py * w + px]: + if py > rmax: + rmax = py + if px > cmax: + cmax = px + if rmax < 0: + for j in range(h * w): + buf[o + j] = False + tw[c] = 1 + th_out[c] = 1 + else: + tw[c] = cmax + 1 + th_out[c] = rmax + 1 + # unique-edge perimeter via sorted int64 keys + Fc = f1 - f0 + if Fc > 0 and V > 0: + keys = np.empty(Fc * 3, dtype=np.int64) + for fi in range(f0, f1): + for j in range(3): + a = faces[fi, j] + b = faces[fi, (j + 1) % 3] + if a < b: + keys[(fi - f0) * 3 + j] = a * V + b + else: + keys[(fi - f0) * 3 + j] = b * V + a + keys = np.sort(keys) + p = 0.0 + for i in range(keys.shape[0]): + if i > 0 and keys[i] == keys[i - 1]: continue - xx = x + i - if xx >= aw: - continue - if atlas[yy, xx]: - ok = False + a = keys[i] // V + v0 + b = keys[i] % V + v0 + dx = rot_uv[a, 0] - rot_uv[b, 0] + dy = rot_uv[a, 1] - rot_uv[b, 1] + p += math.sqrt(dx * dx + dy * dy) + perim[c] = p + + +@_njit(cache=True, boundscheck=False, parallel=True) +def _place_all_jit(buf, boff, stride_w, tw, th, order, start, stop, + atlas, skyline, pool, attempts, sweep_cap, margin, + n_threads, cur_wh, out_x, out_y, out_sw): + """Place charts order[start:stop]; returns the first index NOT processed (== stop when + done, earlier when the atlas must grow — the caller resizes and resumes). The candidate + scan is striped with a (score, index) min-reduction: deterministic for any thread count, + and no thread intrinsics (dynamic globals would defeat cache=True).""" + aw = atlas.shape[1] + ah = atlas.shape[0] + cur_w = cur_wh[0] + cur_h = cur_wh[1] + n_pool = pool.shape[0] + big = np.int64(1) << 62 + nt = n_threads + t_score = np.empty(nt, dtype=np.int64) + t_k = np.empty(nt, dtype=np.int64) + t_x = np.empty(nt, dtype=np.int64) + t_y = np.empty(nt, dtype=np.int64) + t_sw = np.empty(nt, dtype=np.int64) + for oi in range(start, stop): + ci = order[oi] + if cur_h + margin > ah or cur_w + margin > aw: + cur_wh[0] = cur_w + cur_wh[1] = cur_h + return oi + w0 = tw[ci] # unswapped trimmed dims + h0 = th[ci] + W = stride_w[ci] # row stride of the untrimmed block + o = boff[ci] + step = min(w0, h0) // 8 + if step < 1: + step = 1 + cap_step = max(cur_w, cur_h) // sweep_cap + if cap_step > step: + step = cap_step + + poff = (oi * attempts) % (n_pool - attempts + 1) + x_range = cur_w + 1 if cur_w > 0 else 1 + y_range = cur_h + 1 if cur_h > 0 else 1 + # candidate groups per orientation: skyline-flush sweep, y=0 / y=cur_h sweeps, + # x=0 / x=cur_w sweeps; then the shared random pool + nx = max(cur_w, 1) // step + 2 + ny = max(cur_h, 1) // step + 2 + n_det = nx * 3 + ny * 2 + total = n_det * 2 + attempts + for t in range(nt): + t_score[t] = big + t_k[t] = big + for t2 in _prange(nt): + for k in range(t2, total, nt): + x = 0 # int inits and no body-level continue: + y = 0 # parfor lowering types undef-path + swap = 0 # variables as f64 + valid = True + if k < 2 * n_det: + if k >= n_det: + swap = 1 + kk = k - n_det if swap == 1 else k + cw = w0 if swap == 0 else h0 + if kk < nx: # skyline-flush sweep + x = kk * step + if x > cur_w: + valid = False + else: + x_end = x + cw + if x_end > skyline.shape[0]: + x_end = skyline.shape[0] + for xs in range(x, x_end): + if skyline[xs] > y: + y = int(skyline[xs]) + elif kk < 3 * nx: # y=0 and y=cur_h sweeps + kk2 = kk - nx + x = (kk2 % nx) * step + if x > cur_w: + valid = False + elif kk2 >= nx: + y = cur_h + else: # x=0 and x=cur_w sweeps + kk2 = kk - 3 * nx + if kk2 >= 2 * ny: + valid = False + else: + y = (kk2 % ny) * step + if y > cur_h: + valid = False + elif kk2 >= ny: + x = cur_w + else: + r = k - 2 * n_det + x = int(pool[poff + r, 0] % x_range) + y = int(pool[poff + r, 1] % y_range) + swap = int(r & 1) + + if valid: + ch = h0 if swap == 0 else w0 + cw = w0 if swap == 0 else h0 + nw = cur_w if cur_w > x + cw else x + cw + nh = cur_h if cur_h > y + ch else y + ch + ext = nw if nw > nh else nh + score = ext * ext + nw * nh + if score < t_score[t2] or (score == t_score[t2] and k < t_k[t2]): + ok = True + for j in range(ch): + yy = int(y + j) + if yy >= ah: + continue + for i in range(cw): + if swap == 0: + bit = buf[o + j * W + i] + else: + # 90deg rotation: bm_rot[j, i] = bm[h0-1-i, j] + bit = buf[o + (h0 - 1 - i) * W + j] + if not bit: + continue + xx = int(x + i) + if xx >= aw: + continue + if atlas[yy, xx]: + ok = False + break + if not ok: + break + if ok: + t_score[t2] = score + t_k[t2] = k + t_x[t2] = x + t_y[t2] = y + t_sw[t2] = swap + + best_x = -1 + best_y = -1 + best_swap = 0 + bs = big + bk = big + for t in range(nt): + if t_score[t] < bs or (t_score[t] == bs and t_k[t] < bk): + bs = t_score[t] + bk = t_k[t] + best_x = t_x[t] + best_y = t_y[t] + best_swap = t_sw[t] + + if best_x < 0: # fallback: extension corner + best_x = cur_w + best_y = 0 + best_swap = 0 + bh_ = h0 if best_swap == 0 else w0 + bw_ = w0 if best_swap == 0 else h0 + # blit + extents + skyline lift + for j in range(bh_): + for i in range(bw_): + if best_swap == 0: + bit = buf[o + j * W + i] + else: + bit = buf[o + (h0 - 1 - i) * W + j] + if bit: + atlas[best_y + j, best_x + i] = True + if best_x + bw_ > cur_w: + cur_w = best_x + bw_ + if best_y + bh_ > cur_h: + cur_h = best_y + bh_ + for i in range(bw_): + col_x = best_x + i + if col_x >= skyline.shape[0]: + continue + col_top = -1 + for j in range(bh_ - 1, -1, -1): + if best_swap == 0: + bit = buf[o + j * W + i] + else: + bit = buf[o + (h0 - 1 - i) * W + j] + if bit: + col_top = j break - if not ok: - break - if not ok: - continue - best_x = x - best_y = y - best_score = score - best_swap = swap - if x + cw <= cur_w and y + ch <= cur_h: - break - return best_x, best_y, best_score, best_swap - - -def _blit(atlas: np.ndarray, chart: np.ndarray, x: int, y: int) -> None: - ah, aw = atlas.shape - ch, cw = chart.shape - atlas[y: y + ch, x: x + cw] |= chart - - -@dataclass -class _PreparedChart: - chart_id: int - uvs_tex: np.ndarray # [V, 2] in texel coords (rotated, scaled, origin 0) - bitmap: np.ndarray # [h, w] bool, padded - bitmap_rot: np.ndarray # 90° rotated bitmap (for swap_xy placement) - bbox_w: int - bbox_h: int - rotation: float # radians, applied to UVs - s_tex: float # texels per UV unit - perimeter: float # for chart ordering - - -@_njit(cache=True, boundscheck=False) -def _chart_perimeter_jit(uvs: np.ndarray, faces: np.ndarray, V: int) -> float: - """Sum unique-edge lengths via sorted int64 edge keys.""" - F = faces.shape[0] - keys = np.empty(F * 3, dtype=np.int64) - for fi in range(F): - for j in range(3): - a = faces[fi, j] - b = faces[fi, (j + 1) % 3] - if a < b: - keys[fi * 3 + j] = a * V + b - else: - keys[fi * 3 + j] = b * V + a - keys = np.sort(keys) - p = 0.0 - for i in range(keys.shape[0]): - if i > 0 and keys[i] == keys[i - 1]: - continue - a = keys[i] // V - b = keys[i] % V - dx = uvs[a, 0] - uvs[b, 0] - dy = uvs[a, 1] - uvs[b, 1] - p += math.sqrt(dx * dx + dy * dy) - return p - - -def _chart_perimeter(uvs: np.ndarray, faces: np.ndarray) -> float: - V = int(faces.max()) + 1 if faces.size else 0 - return float(_chart_perimeter_jit(uvs.astype(np.float64), faces.astype(np.int64), V)) + if col_top >= 0: + nh2 = best_y + col_top + 1 + if nh2 > skyline[col_x]: + skyline[col_x] = nh2 + out_x[ci] = best_x + out_y[ci] = best_y + out_sw[ci] = best_swap + cur_wh[0] = cur_w + cur_wh[1] = cur_h + return stop # Torch fallback (used when numba is unavailable; runs on GPU if present) def _dilate_local(x: Tensor, p: int) -> Tensor: - """4-connectivity dilation by p, applied per-image over a batch of (cnt,g,g) bitmaps. - Matches the old per-chart _dilate_torch; dilation distributes over union so per-triangle - dilation OR-scattered equals dilating the assembled chart bitmap.""" + """4-connectivity dilation by p over a batch of (cnt,g,g) bitmaps. Dilation distributes + over union, so dilating per-triangle then OR-scattering equals dilating the chart.""" for _ in range(p): y = x.clone() y[:, 1:, :] |= x[:, :-1, :] @@ -380,10 +430,8 @@ def _dilate_local(x: Tensor, p: int) -> Tensor: def _raster_all_torch(uvs_tex_pad, faces_pad, fmask, bw_t, bh_t, padding, device): - """Batched rasterize EVERY chart at once into one flat bool buffer, replacing the per-chart - loop. Returns (buf, cbase) where buf[cbase[i]:cbase[i+1]].view(bh,bw) is chart i's [y,x] bitmap. - Triangles are bucketed by next-pow2 bbox size so each batch's local grid stays tiny (bounded - memory) while collapsing ~N chart rasters into a handful of kernels.""" + """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] bwL, bhL = bw_t.long(), bh_t.long() @@ -450,66 +498,99 @@ def _raster_all_torch(uvs_tex_pad, faces_pad, fmask, bw_t, bh_t, padding, device return buf, cbase -def _build_candidates_gpu(sky_t, cur_w, cur_h, bw0, bw1, step, rand_n, gen, device): - """Skyline-flush + edge-sweep + random candidate (x,y) positions per orientation, built on the - GPU. Returns (cand0, cand1). Random samples find tight pockets the deterministic grid misses.""" - xs = torch.arange(0, max(cur_w, 1) + 1, step, device=device) - ys = torch.arange(0, max(cur_h, 1) + 1, step, device=device) - # edge-sweep candidates are orientation-independent: build once, shared by both orientations - common = [torch.stack([xs, torch.full_like(xs, yf)], 1) for yf in (0, cur_h)] - common += [torch.stack([torch.full_like(ys, xf), ys], 1) for xf in (0, cur_w)] - common = torch.cat(common, 0) - out = [] - for cw in (bw0, bw1): # skyline-flush + random differ - if cw > 0 and sky_t.shape[0] >= cw: - wmax = sky_t.unfold(0, cw, 1).amax(1)[xs.clamp(max=max(sky_t.shape[0] - cw, 0))] - else: - wmax = torch.zeros_like(xs) - parts = [torch.stack([xs, wmax], 1), common] - if rand_n > 0: # distinct draws keep density - rx = torch.randint(0, max(cur_w, 1) + 1, (rand_n,), generator=gen, device=device) - ry = torch.randint(0, max(cur_h, 1) + 1, (rand_n,), generator=gen, device=device) - parts.append(torch.stack([rx, ry], 1)) - out.append(torch.cat(parts, 0)) - return out[0], out[1] +def _build_candidates_gpu(sky_t, ar, cur_w, cur_h, bw0, bw1, step, rand01, device): + """Candidate (x, y) positions as a (2, M, 2) tensor (dim 0 = orientation). The first + n_sky rows per orientation are skyline-flush and collision-free by construction. + rand01 is (2, rand_n, 2) pre-drawn uniforms; ar a preallocated arange.""" + hi_x = max(cur_w, 1) + 1 + hi_y = max(cur_h, 1) + 1 + xs = ar[0:hi_x:step] + ys = ar[0:hi_y:step] + n_sky = (hi_x + step - 1) // step + zx = torch.zeros_like(xs) + zy = torch.zeros_like(ys) + common = torch.cat([ + torch.stack([xs, zx], 1), torch.stack([xs, zx + cur_h], 1), + torch.stack([zy, ys], 1), torch.stack([zy + cur_w, ys], 1)]) + wm = [] + for cw in (bw0, bw1): + span = (n_sky - 1) * step + cw + wm.append(max_pool1d(sky_t[:span].view(1, 1, -1).float(), kernel_size=cw, + stride=step).view(-1)) + sky = torch.stack([torch.stack([xs, wm[0].long()], 1), + torch.stack([xs, wm[1].long()], 1)]) + lim = torch.tensor([hi_x, hi_y], dtype=rand01.dtype, device=device) + rnd = (rand01 * lim).long() + return torch.cat([sky, common.expand(2, -1, -1), rnd], 1), n_sky -def _col_top(b: Tensor) -> Tensor: - """Topmost True row index per column of a bool bitmap (h,w); -1 for empty columns.""" - h = b.shape[0] - rows = torch.arange(h, device=b.device)[:, None] - return torch.where(b, rows, torch.full_like(rows.expand_as(b), -1)).amax(0) +def _best_placement_torch(atlas, pix0, dim0, dim1, cands, n_sky, cur_w, cur_h, device): + """Lowest-score non-colliding placement as a (3,) int tensor [x, y, swap]. The best + skyline candidate bounds the score; only strictly better candidates are pixel-tested.""" + m = cands.shape[1] + chw = torch.tensor([[dim0[0], dim0[1]], [dim1[0], dim1[1]]], device=device) + nw = torch.clamp(cands[..., 0] + chw[:, 1:], min=cur_w) # (2,M) + nh = torch.clamp(cands[..., 1] + chw[:, :1], min=cur_h) + ext = torch.maximum(nw, nh) + sc = ext * ext + nw * nh + js = sc[:, :n_sky].reshape(-1).argmin() # best skyline candidate + sky_o = js // n_sky + s_star = sc[:, :n_sky].reshape(-1)[js] + sky = torch.cat([cands[sky_o, js % n_sky], sky_o.reshape(1)]) + cflat = cands.reshape(-1, 2) + surv = (sc.reshape(-1) < s_star).nonzero(as_tuple=True)[0] # compact once + total = surv.shape[0] + if total == 0: + return sky + k = pix0.shape[0] + if k == 0: # empty chart: anywhere free + j = surv[sc.reshape(-1)[surv].argmin()] + return torch.cat([cflat[j], (j // m).reshape(1)]) + ordr = surv[torch.argsort(sc.reshape(-1)[surv], stable=True)] -def _best_placement_torch(atlas, pix0, dim0, pix1, dim1, cand0, cand1, cur_w, cur_h, device): - """Lowest-score non-colliding candidate as a (3,) int tensor [x, y, swap] (x=-1 if none). - Collision tests only each bitmap's True-pixel offsets (pix), not the full window. Fully on-GPU; - the caller does the single sync (.tolist()).""" - INF = 1 << 60 + # flattened-index collision test: one int32 gather index instead of two int64 rows/cols + aw = atlas.shape[1] + idt = torch.int32 if atlas.numel() < (1 << 31) else torch.long + lin0 = (pix0[:, 0] * aw + pix0[:, 1]).to(idt) # (y, x) + lin1 = (pix0[:, 1] * aw + (dim0[0] - 1 - pix0[:, 0])).to(idt) # rotated: (x, h-1-y) + linp = torch.stack([lin0, lin1]) + aflat = atlas.view(-1) + og = (ordr >= m).long() + base = (cflat[ordr, 1] * aw + cflat[ordr, 0]).to(idt) - def best(cand, pix, dim): # -> (score, x, y) 0-d tensors - ch, cw = dim - cx, cy = cand[:, 0], cand[:, 1] - coll = atlas[cy[:, None] + pix[:, 0][None, :], # (M,k) True-pixel gather - cx[:, None] + pix[:, 1][None, :]].any(dim=1) - nw = torch.clamp(cx + cw, min=cur_w) - nh = torch.clamp(cy + ch, min=cur_h) - ext = torch.maximum(nw, nh) - score = torch.where(coll, torch.full_like(nw, INF), ext * ext + nw * nh) - j = score.argmin() - return score[j], cx[j], cy[j] + # prescreen survivors on ~128 strided pixels: a sampled hit proves collision, so only + # subsample-clean candidates need the exact test + stride = (k + 127) // 128 + linp_sub = linp[:, ::stride].contiguous() + maybe = ~aflat[base[:, None] + linp_sub[og]].any(1) + passers = maybe.nonzero(as_tuple=True)[0] # ascending = score-sorted + npass = passers.shape[0] + if npass == 0: + return sky + if stride == 1: # prescreen was already exact + j = ordr[passers[0]] + return torch.cat([cflat[j], (j // m).reshape(1)]) - s0, x0, y0 = best(cand0, pix0, dim0) - s1, x1, y1 = best(cand1, pix1, dim1) - take0 = s0 <= s1 - bsc = torch.where(take0, s0, s1) - pick = torch.stack([torch.where(take0, x0, x1), torch.where(take0, y0, y1), - torch.where(take0, x0.new_zeros(()), x0.new_ones(()))]) - return torch.where(bsc < INF, pick, torch.tensor([-1, -1, 0], device=device)) + budget = 1 << 22 # pixel-tests per chunk + start = 0 + while start < npass: + take = max(1, budget // k) + pi = passers[start:start + take] + free = ~aflat[base[pi][:, None] + linp[og[pi]]].any(1) # (t,k) True-pixel gather + # single host read per chunk: whether a free hit exists and where + has, first = torch.stack([free.any().long(), free.long().argmax()]).tolist() + if has: + j = ordr[pi[first]] # lowest score: sorted order + return torch.cat([cflat[j], (j // m).reshape(1)]) + start += take + budget = min(budget * 4, 1 << 25) + return sky def _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces, - texels_per_unit, padding_texels): + 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) if n == 0: @@ -563,47 +644,46 @@ def _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces, # one sync: pull all per-chart scalars thetas = ang[ti].cpu().tolist() scales = scale.cpu().tolist() - bws = bw_t.cpu().tolist() - bhs = bh_t.cpu().tolist() - # ---- Prepare pass 2: rasterize ALL charts at once, then trim each bitmap to its bounds ---- + # ---- 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) - cb = cbase.cpu().tolist() - raw, bnd = [], [] - for i in range(n): - bm = buf[cb[i]:cb[i + 1]].view(bhs[i], bws[i]) - raw.append(bm) - rr = torch.arange(bm.shape[0], device=device) - cc = torch.arange(bm.shape[1], device=device) - rmax = torch.where(bm.any(1), rr, rr.new_full((), -1)).amax() # last occupied row / col (-1 if empty) - cmax = torch.where(bm.any(0), cc, cc.new_full((), -1)).amax() - bnd.append(torch.stack([rmax, cmax])) - bnd_cpu = torch.stack(bnd).cpu().tolist() # one sync for all trim bounds - # per-chart True-pixel offsets (sparse collision/blit), dims, col-tops (all kept on GPU) - pix_l, pixr_l, dim_l, dimr_l, bm_h = [], [], [], [], [] - col_tops, col_tops_rot = [], [] - for i in range(n): - rm, cm = bnd_cpu[i] - bm = (raw[i][:rm + 1, :cm + 1].contiguous() if rm >= 0 and cm >= 0 - else torch.zeros((1, 1), dtype=torch.bool, device=device)) - bm_rot = torch.flip(bm.t(), dims=[1]).contiguous() - pix_l.append(bm.nonzero()) - pixr_l.append(bm_rot.nonzero()) - dim_l.append((int(bm.shape[0]), int(bm.shape[1]))) - dimr_l.append((int(bm_rot.shape[0]), int(bm_rot.shape[1]))) - col_tops.append(_col_top(bm)) - col_tops_rot.append(_col_top(bm_rot)) - bm_h.append(int(bm.shape[0])) - wmax = max(d[1] for d in dim_l + dimr_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) - for i in range(n): - ct_pad[i, :col_tops[i].shape[0]] = col_tops[i] - ctr_pad[i, :col_tops_rot[i].shape[0]] = col_tops_rot[i] - del raw + # nonzero over the flat buffer is ascending, so pixels come out grouped by chart + nz = buf.nonzero(as_tuple=True)[0] + del buf + cid = torch.searchsorted(cbase, nz, right=True) - 1 + bwl = bw_t.long() + local = nz - cbase[cid] + py = local // bwl[cid] + px = local - py * bwl[cid] + del nz, local + counts = torch.bincount(cid, minlength=n) + rmax = torch.full((n,), -1, dtype=torch.long, device=device) + cmax = torch.full((n,), -1, dtype=torch.long, device=device) + rmax.scatter_reduce_(0, cid, py, reduce="amax") + cmax.scatter_reduce_(0, cid, px, reduce="amax") + ht = (rmax + 1).clamp_min(1) # trimmed bitmap dims (1x1 when empty) + wt = (cmax + 1).clamp_min(1) + pix_all = torch.stack([py, px], 1) # True-pixel (row, col) offsets, sparse + pixr_all = torch.stack([px, rmax[cid] - py], 1) # 90deg rotation: (y, x) -> (x, h-1-y) + meta = torch.stack([ht, wt, counts.cumsum(0)], 1).cpu().tolist() # one sync for all charts + dim_l = [(m[0], m[1]) for m in meta] + dimr_l = [(w, h) for (h, w) in dim_l] + offs = [0] + [m[2] for m in meta] + 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)] - # ---- Placement: skyline bin-pack on GPU (1 sync/chart for the chosen position) ---- + # 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 ---- order = sorted(range(n), key=lambda i: -(dim_l[i][0] * dim_l[i][1])) # biggest bitmap first max_b = max(max(d) for d in dim_l) margin = max_b + 8 @@ -611,12 +691,17 @@ def _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces, cap = side_guess + margin atlas = torch.zeros((cap, cap), dtype=torch.bool, device=device) sky_t = torch.zeros(cap, dtype=torch.long, device=device) + ar = torch.arange(cap + 1, device=device) cur_w = cur_h = 0 placements = [None] * n - gen = torch.Generator(device=device).manual_seed(0) - rand_n = 512 # random samples per orientation + gen = torch.Generator(device=device).manual_seed(rng_seed) + rand_n = min(512, attempts) # random samples per orientation + # no _SWEEP_CAP here: the skyline-bound pruning depends on the dense sweep + rand01 = torch.rand(n, 2, rand_n, 2, generator=gen, device=device) # all draws upfront - for ci in order: + for t_i, ci in enumerate(order): + if progress_callback is not None and (t_i & 255) == 0: + progress_callback(n + t_i, 2 * n) if cur_h + margin > atlas.shape[0] or cur_w + margin > atlas.shape[1]: ns = max(atlas.shape[0], cur_h + margin, cur_w + margin) na = torch.zeros((ns, ns), dtype=torch.bool, device=device) @@ -625,11 +710,14 @@ def _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces, nsk = torch.zeros(ns, dtype=torch.long, device=device) nsk[:sky_t.shape[0]] = sky_t sky_t = nsk + ar = torch.arange(ns + 1, device=device) dim, dimr = dim_l[ci], dimr_l[ci] step = max(1, min(dim[0], dim[1]) // 8) - cand0, cand1 = _build_candidates_gpu(sky_t, cur_w, cur_h, dim[1], dimr[1], step, rand_n, gen, device) - res = _best_placement_torch(atlas, pix_l[ci], dim, pixr_l[ci], dimr, cand0, cand1, cur_w, cur_h, device) - bx, by, swap = (int(v) for v in res.tolist()) # the one sync/chart + cands, n_sky = _build_candidates_gpu( + sky_t, ar, cur_w, cur_h, dim[1], dimr[1], step, rand01[t_i], device) + res = _best_placement_torch(atlas, pix_l[ci], dim, dimr, + cands, n_sky, cur_w, cur_h, device) + bx, by, swap = (int(v) for v in res.tolist()) if bx < 0: bx, by, swap = cur_w, 0, 0 pix = pixr_l[ci] if swap else pix_l[ci] @@ -638,189 +726,136 @@ def _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces, 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 - ix = torch.arange(bx, bx + bw_, device=device) + 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)), scale=scales[ci], rotation=thetas[ci], swap_xy=bool(swap), - chart_h=float(bm_h[ci])) + chart_h=float(dim_l[ci][0])) return placements, cur_w, cur_h -def pack_bitmap( - chart_uvs: List[Tensor], - chart_3d_areas: List[float], - chart_uv_areas: List[float], - chart_faces: List[Tensor], +def pack_bitmap_concat( + uvs_cat: np.ndarray, # (sumV, 2) per-chart concatenated UVs + uv_offsets: np.ndarray, # (n+1,) + faces_cat: np.ndarray, # (sumF, 3) local vert ids per chart + face_offsets: np.ndarray, # (n+1,) + chart_3d_areas: np.ndarray, + chart_uv_areas: np.ndarray, texels_per_unit: float = 256.0, padding_texels: int = 2, attempts: int = 4096, rng_seed: int = 0, -) -> Tuple[List[ChartPlacement], int, int]: - """Rasterize-and-place packer. Returns (placements, atlas_w, atlas_h).""" - n = len(chart_uvs) + progress_callback=None, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int]: + """Rasterize-and-place packer over concatenated chart arrays (no per-chart python). + Returns (x, y, swap, rotation, scale, chart_h, atlas_w, atlas_h) with one entry per chart. + progress_callback(done, total) is invoked periodically; total is 2*n_charts.""" + n = int(uv_offsets.shape[0]) - 1 + empty = np.zeros(n, dtype=np.int64) if n == 0: - return [], 1, 1 + return empty, empty, empty, empty.astype(np.float64), empty.astype(np.float64), empty, 1, 1 if not _HAVE_NUMBA_PACK: - return _pack_bitmap_torch(chart_uvs, chart_3d_areas, chart_uv_areas, - chart_faces, texels_per_unit, padding_texels) + 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, + 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) + sw = np.array([1 if p.swap_xy else 0 for p in placements], dtype=np.int64) + th = np.array([p.rotation for p in placements], dtype=np.float64) + sc = np.array([p.scale for p in placements], dtype=np.float64) + chh = np.array([p.chart_h for p in placements], dtype=np.int64) + return px, py, sw, th, sc, chh, w, h + uvs64 = np.ascontiguousarray(uvs_cat, dtype=np.float64) + faces64 = np.ascontiguousarray(faces_cat, dtype=np.int64) + uv_off = np.ascontiguousarray(uv_offsets, dtype=np.int64) + f_off = np.ascontiguousarray(face_offsets, dtype=np.int64) + a3 = np.ascontiguousarray(chart_3d_areas, dtype=np.float64) + auv = np.ascontiguousarray(chart_uv_areas, dtype=np.float64) + + theta = np.zeros(n, dtype=np.float64) + scale = np.zeros(n, dtype=np.float64) + bw = np.zeros(n, dtype=np.int64) + bh = np.zeros(n, dtype=np.int64) + rot_uv = np.empty_like(uvs64) + _prepare_dims_jit(uvs64, uv_off, a3, auv, float(texels_per_unit), int(padding_texels), + theta, scale, bw, bh, rot_uv) + boff = np.zeros(n + 1, dtype=np.int64) + np.cumsum(bw * bh, out=boff[1:]) + buf = np.zeros(int(boff[-1]), dtype=np.bool_) + tw = np.zeros(n, dtype=np.int64) + th_arr = np.zeros(n, dtype=np.int64) + perim = np.zeros(n, dtype=np.float64) + _raster_all_jit(rot_uv, uv_off, faces64, f_off, bw, bh, boff, buf, + int(padding_texels), tw, th_arr, perim) + if progress_callback is not None: + progress_callback(n, 2 * n) + + order = np.argsort(-perim, kind="stable") + max_b = int(max(int(tw.max()), int(th_arr.max()))) + margin = max_b + 8 + side_guess = int(math.sqrt(float((tw * th_arr).sum()))) * 2 + 16 + cap = side_guess + margin + atlas = np.zeros((cap, cap), dtype=np.bool_) + skyline = np.zeros(cap, dtype=np.int64) rng = np.random.default_rng(rng_seed) - prepared: List[_PreparedChart] = [] - skyline_cap = 4096 - skyline = np.zeros(skyline_cap, dtype=np.int64) - - for i, (uvs_t, area_3d, area_uv, faces_t) in enumerate( - zip(chart_uvs, chart_3d_areas, chart_uv_areas, chart_faces) - ): - uvs = uvs_t.detach().cpu().numpy().astype(np.float64) - faces = faces_t.detach().cpu().numpy() - - theta = _best_rotation(uvs) - rotated = _rotate_xy(uvs, theta) - scale = math.sqrt(max(area_3d, 1e-12) / max(area_uv, 1e-12)) * texels_per_unit - # Cap per-chart bbox to 4x nominal so a degenerate chart can't span the atlas. - nominal_side = math.sqrt(max(area_3d, 1e-12)) * float(texels_per_unit) - max_bbox_texels = max(8.0, 4.0 * nominal_side) - bbox_uv = (rotated.max(axis=0) - rotated.min(axis=0)) - bbox_uv_max = float(max(bbox_uv[0], bbox_uv[1], 1e-12)) - if scale * bbox_uv_max > max_bbox_texels: - scale = max_bbox_texels / bbox_uv_max - uvs_tex = rotated * scale - uvs_tex = uvs_tex - uvs_tex.min(axis=0) - bbox_w = int(math.ceil(uvs_tex[:, 0].max())) + padding_texels + 1 - bbox_h = int(math.ceil(uvs_tex[:, 1].max())) + padding_texels + 1 - - bm = _rasterize_chart(uvs_tex, faces, bbox_w, bbox_h, padding_texels) - nz_rows = np.where(bm.any(axis=1))[0] - nz_cols = np.where(bm.any(axis=0))[0] - if nz_rows.size == 0 or nz_cols.size == 0: - bm = np.zeros((1, 1), dtype=bool) - bbox_h, bbox_w = 1, 1 - else: - bm = bm[: nz_rows[-1] + 1, : nz_cols[-1] + 1] - bbox_h, bbox_w = bm.shape - # True 90 deg rotation; plain transpose would mirror and flip winding. - bm_rot = bm.T[:, ::-1].copy() - - perim = _chart_perimeter(uvs_tex, faces) - prepared.append( - _PreparedChart( - chart_id=i, - uvs_tex=uvs_tex, - bitmap=bm, - bitmap_rot=bm_rot, - bbox_w=bbox_w, - bbox_h=bbox_h, - rotation=theta, - s_tex=scale, - perimeter=perim, - ) - ) - - order = sorted(range(n), key=lambda i: -prepared[i].perimeter) - - total_area = sum(p.bbox_w * p.bbox_h for p in prepared) - side_guess = int(math.sqrt(total_area) * 2) + 16 - atlas = np.zeros((side_guess, side_guess), dtype=bool) - cur_w = 0 - cur_h = 0 - - placements: List[ChartPlacement] = [None] * n # type: ignore - - for ci in order: - p = prepared[ci] - - step = max(1, min(p.bbox_w, p.bbox_h) // 8) - det_arr = _build_candidates_jit( - skyline, cur_w, cur_h, - p.bitmap.shape[1], p.bitmap.shape[0], - p.bitmap_rot.shape[1], p.bitmap_rot.shape[0], - step, - ) - - x_range = max(cur_w + 1, 1) - y_range = max(cur_h + 1, 1) - rand_x = rng.integers(0, x_range, size=attempts).astype(np.int64) - rand_y = rng.integers(0, y_range, size=attempts).astype(np.int64) - rand_swap = (np.arange(attempts) & 1).astype(np.int64) - rand_arr = np.stack([rand_x, rand_y, rand_swap], axis=1) - candidates = np.concatenate([det_arr, rand_arr], axis=0) if det_arr.size else rand_arr - - best_x, best_y, best_score_int, best_swap_int = _best_placement_jit( - atlas, p.bitmap, p.bitmap_rot, candidates, cur_w, cur_h, - ) - best_swap = bool(best_swap_int) - - if best_x >= 0: - bm_b = p.bitmap_rot if best_swap else p.bitmap - need_h = max(cur_h, best_y + bm_b.shape[0]) - need_w = max(cur_w, best_x + bm_b.shape[1]) - if atlas.shape[0] < need_h or atlas.shape[1] < need_w: - target_h = max(atlas.shape[0], need_h, side_guess) - target_w = max(atlas.shape[1], need_w, side_guess) - new_atlas = np.zeros((target_h, target_w), dtype=bool) - new_atlas[: atlas.shape[0], : atlas.shape[1]] = atlas - atlas = new_atlas - - if best_x < 0: - # Fallback: place at extension corner. - best_x, best_y = cur_w, 0 - best_swap = False - bm = p.bitmap - need_h = max(cur_h, best_y + bm.shape[0]) - need_w = max(cur_w, best_x + bm.shape[1]) - if atlas.shape[0] < need_h or atlas.shape[1] < need_w: - target_h = max(atlas.shape[0], need_h) - target_w = max(atlas.shape[1], need_w) - new_atlas = np.zeros((target_h, target_w), dtype=bool) - new_atlas[: atlas.shape[0], : atlas.shape[1]] = atlas - atlas = new_atlas - - bm = p.bitmap_rot if best_swap else p.bitmap - _blit(atlas, bm, best_x, best_y) - cur_w = max(cur_w, best_x + bm.shape[1]) - cur_h = max(cur_h, best_y + bm.shape[0]) - if cur_w + 1 > skyline.shape[0]: - new_sky = np.zeros(max(skyline.shape[0] * 2, cur_w + 1), dtype=np.int64) - new_sky[: skyline.shape[0]] = skyline - skyline = new_sky - _update_skyline_jit(skyline, bm, best_x, best_y) - - placements[ci] = ChartPlacement( - chart_id=ci, - offset=(float(best_x), float(best_y)), - scale=p.s_tex, - rotation=p.rotation, - swap_xy=best_swap, - chart_h=float(p.bitmap.shape[0]), - ) - - return placements, cur_w, cur_h + # shared random pool, sliced at a rotating offset per chart + pool = rng.integers(0, 1 << 31, size=(attempts * 8, 2)).astype(np.int64) + out_x = np.full(n, -1, dtype=np.int64) + out_y = np.full(n, -1, dtype=np.int64) + out_sw = np.zeros(n, dtype=np.int64) + cur_wh = np.zeros(2, dtype=np.int64) + start = 0 + while start < n: + stop = min(n, start + 1024) + nxt = _place_all_jit(buf, boff, bw, tw, th_arr, order, start, stop, + atlas, skyline, pool, int(attempts), int(_SWEEP_CAP), + int(margin), int(_nb_threads()), cur_wh, out_x, out_y, out_sw) + if nxt < stop: # atlas must grow before this chart fits + ns = max(atlas.shape[0], int(cur_wh[1]) + margin, int(cur_wh[0]) + margin) + na = np.zeros((ns, ns), dtype=np.bool_) + na[:atlas.shape[0], :atlas.shape[1]] = atlas + atlas = na + nsk = np.zeros(ns, dtype=np.int64) + nsk[:skyline.shape[0]] = skyline + skyline = nsk + start = nxt + if progress_callback is not None: + progress_callback(n + start, 2 * n) + return out_x, out_y, out_sw, theta, scale, th_arr, int(cur_wh[0]), int(cur_wh[1]) -def apply_placements( - chart_uvs: List[Tensor], placements: List[ChartPlacement], atlas_w: int, atlas_h: int -) -> List[Tensor]: - """Apply per-chart (rotation, scale, swap_xy, offset) and normalize by the larger atlas side (shared scale keeps texel density uniform).""" - out: List[Tensor] = [] +def apply_placements_concat( + uvs_cat: np.ndarray, uv_offsets: np.ndarray, + px: np.ndarray, py: np.ndarray, sw: np.ndarray, + theta: np.ndarray, scale: np.ndarray, chart_h: np.ndarray, + atlas_w: int, atlas_h: int, +) -> np.ndarray: + """apply_placements over concatenated charts, fully vectorized. Returns (sumV, 2) float32.""" + n = int(uv_offsets.shape[0]) - 1 side = float(max(atlas_w, atlas_h, 1)) - for uvs, p in zip(chart_uvs, placements): - device = uvs.device - dtype = uvs.dtype - uvs_np = uvs.detach().cpu().numpy().astype(np.float64) - if p.rotation != 0.0: - uvs_np = _rotate_xy(uvs_np, p.rotation) - uvs_np = uvs_np - uvs_np.min(axis=0) - uvs_np = uvs_np * p.scale - if p.swap_xy: - # 90 deg rotation matching bm.T[:, ::-1]: (u, v) -> (chart_h - v, u). - u_old = uvs_np[:, 0].copy() - uvs_np[:, 0] = p.chart_h - uvs_np[:, 1] - uvs_np[:, 1] = u_old - uvs_np[:, 0] += p.offset[0] - uvs_np[:, 1] += p.offset[1] - uvs_np /= side - # Clamp into [0,1]; slivers can stick sub-texel past the tracked extent. - np.clip(uvs_np, 0.0, 1.0, out=uvs_np) - out.append(torch.from_numpy(uvs_np).to(device=device, dtype=dtype)) - return out + cov = np.repeat(np.arange(n), np.diff(uv_offsets)) + u_in = uvs_cat[:, 0].astype(np.float64) + v_in = uvs_cat[:, 1].astype(np.float64) + c = np.cos(theta)[cov] + s = np.sin(theta)[cov] + u = u_in * c - v_in * s + v = u_in * s + v_in * c + umin = np.full(n, np.inf) + vmin = np.full(n, np.inf) + np.minimum.at(umin, cov, u) + np.minimum.at(vmin, cov, v) + u = (u - umin[cov]) * scale[cov] + v = (v - vmin[cov]) * scale[cov] + swv = sw[cov].astype(bool) + # 90 deg rotation matching the rotated-bitmap access: (u, v) -> (chart_h - v, u) + u2 = np.where(swv, chart_h[cov] - v, u) + px[cov] + v2 = np.where(swv, u, v) + py[cov] + out = np.stack([u2, v2], axis=1) / side + np.clip(out, 0.0, 1.0, out=out) # slivers can stick sub-texel past extents + return out.astype(np.float32) diff --git a/comfy_extras/mesh3d/uv_unwrap/parameterize.py b/comfy_extras/mesh3d/uv_unwrap/parameterize.py index 8b59575a0..cbf6fa689 100644 --- a/comfy_extras/mesh3d/uv_unwrap/parameterize.py +++ b/comfy_extras/mesh3d/uv_unwrap/parameterize.py @@ -12,6 +12,8 @@ from torch import Tensor 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: """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) -def _stretch_metrics(verts_3d: np.ndarray, uvs: np.ndarray, faces: np.ndarray) -> Tuple[float, float, int, int]: - """Sander's stretch metric. Returns (rms, max, n_flipped, n_zero_area).""" - p = verts_3d[faces] +def ortho_project_concat(verts: np.ndarray, chart_of_vert: np.ndarray, n_charts: int) -> np.ndarray: + """_ortho_project for every chart at once over concatenated per-chart vertices.""" + 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] - parametric_area = 0.5 * ( + pa_signed = 0.5 * ( (t[:, 1, 1] - t[:, 0, 1]) * (t[:, 2, 0] - t[:, 0, 0]) - - (t[:, 2, 1] - t[:, 0, 1]) * (t[:, 1, 0] - t[:, 0, 0]) - ) - n_flipped = int((parametric_area < -1e-12).sum()) - n_zero = int((np.abs(parametric_area) < 1e-12).sum()) - pa = np.abs(parametric_area).clip(min=1e-20) - geom_area = 0.5 * np.linalg.norm( - np.cross(p[:, 1] - p[:, 0], p[:, 2] - p[:, 0]), axis=1 - ) - keep = (geom_area > 1e-12) & (np.abs(parametric_area) > 1e-12) - 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] + - (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_zero = np.bincount(chart_of_face[np.abs(pa_signed) < 1e-12], minlength=n_charts) + pa = np.abs(pa_signed).clip(min=1e-20) + ga = 0.5 * np.linalg.norm(np.cross(p[:, 1] - p[:, 0], p[:, 2] - p[:, 0]), axis=1) + keep = (ga > 1e-12) & (np.abs(pa_signed) > 1e-12) + t1, s1 = t[:, 0, 0], t[:, 0, 1] + t2, s2 = t[:, 1, 0], t[:, 1, 1] + t3, s3 = t[:, 2, 0], t[:, 2, 1] inv_2pa = 1.0 / (2.0 * pa) - Ss = ( - p[:, 0] * (t2 - t3)[:, None] - + p[:, 1] * (t3 - t1)[:, None] - + p[:, 2] * (t1 - t2)[:, None] - ) * inv_2pa[:, None] - St = ( - p[:, 0] * (s3 - s2)[:, None] - + p[:, 1] * (s1 - s3)[:, None] - + p[:, 2] * (s2 - s1)[:, None] - ) * inv_2pa[:, None] + Ss = (p[:, 0] * (t2 - t3)[:, None] + p[:, 1] * (t3 - t1)[:, None] + + p[:, 2] * (t1 - t2)[:, 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) bb = (Ss * 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))) rms_sq = (a + c) * 0.5 - rms_stretch_sq_sum = float((rms_sq[keep] * geom_area[keep]).sum()) - total_geom = float(geom_area[keep].sum()) - total_param = float(pa[keep].sum()) - if total_geom <= 0.0: - return float("inf"), float("inf"), n_flipped, n_zero - norm_factor = np.sqrt(total_param / total_geom) - rms_stretch = float(np.sqrt(rms_stretch_sq_sum / total_geom)) * norm_factor - max_stretch = float(np.sqrt(sigma2_sq[keep].max())) * norm_factor - return rms_stretch, max_stretch, n_flipped, n_zero + cf = chart_of_face[keep] + tg = np.bincount(cf, weights=ga[keep], minlength=n_charts) + tp = np.bincount(cf, weights=pa[keep], minlength=n_charts) + rs = np.bincount(cf, weights=(rms_sq * ga)[keep], minlength=n_charts) + smax = np.zeros(n_charts, dtype=np.float64) + np.maximum.at(smax, cf, sigma2_sq[keep]) + ok = tg > 0.0 + tg_safe = np.where(ok, tg, 1.0) + 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( @@ -181,36 +387,6 @@ def _uv_boundary_self_intersects( 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( verts_3d: np.ndarray, faces: np.ndarray ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: diff --git a/comfy_extras/mesh3d/uv_unwrap/segment.py b/comfy_extras/mesh3d/uv_unwrap/segment.py index 1983ccca8..df4ba8ee1 100644 --- a/comfy_extras/mesh3d/uv_unwrap/segment.py +++ b/comfy_extras/mesh3d/uv_unwrap/segment.py @@ -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 typing import List, Tuple +from typing import Tuple -import numpy as np import torch from torch import Tensor - -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 tqdm import tqdm from .mesh import MeshData, face_edge_lengths @@ -28,129 +17,63 @@ DEFAULT_MAX_COST = 2.0 NORMAL_DEVIATION_HARD_CUTOFF = 0.707 # ~75° -@njit(cache=True, fastmath=False) -def _cost_grow_iter_jit( - face_chart: np.ndarray, face_face: np.ndarray, face_normal: np.ndarray, - face_area: np.ndarray, face_edge_len: np.ndarray, - chart_basis: np.ndarray, chart_normal_sum: np.ndarray, - chart_area: np.ndarray, chart_perim: np.ndarray, - nd_cutoff: float, max_cost: float, - w_nd: float, w_round: float, w_straight: float, -): - """One grow iter: each unassigned face joins its lowest-cost adjacent chart if cost= 0, face_chart[nb.clamp_min(0)], nb.new_full((), -1)) + valid = nbc >= 0 + d = (fn[u][:, None, :] * basis[nbc.clamp_min(0)]).sum(-1) + nd = (1.0 - d).clamp(0.0, 1.0) + 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): - if face_chart[f] != -1: - continue - nx = face_normal[f, 0] - ny = face_normal[f, 1] - nz = face_normal[f, 2] - af = face_area[f] - for e0 in range(3): - nb0 = face_face[f, e0] - if nb0 < 0: - continue - c = face_chart[nb0] - if c < 0: - continue - d = (nx * chart_basis[c, 0] + ny * chart_basis[c, 1] + nz * chart_basis[c, 2]) - nd = np.float32(1.0) - d - if nd > np.float32(1.0): - nd = np.float32(1.0) - if nd < np.float32(0.0): - nd = np.float32(0.0) - if nd >= nd_cutoff: - continue - 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) + f_acc = u[acc] + c_acc = nbc.gather(1, best_j[:, None]).squeeze(1)[acc] + nbc_old = nbc[acc] # neighbor charts before this commit + face_chart[f_acc] = c_acc + nb_acc = nb[acc] + nbs_acc = nb_acc.clamp_min(0) + nbc_post = torch.where(nb_acc >= 0, face_chart[nbs_acc], nb_acc.new_full((), -1)) + # frontier update: committed faces leave; their still-unassigned neighbors enter + frontier[f_acc] = False + grow_nb = nbs_acc[(nb_acc >= 0) & (nbc_post < 0)] + frontier[grow_nb] = True + el_acc = el[acc] + cx = c_acc[:, None] + dper = torch.where(nbc_old == cx, -el_acc, # was member: edge turns interior + torch.where(nbc_post == cx, torch.zeros_like(el_acc), # co-committer + el_acc)).sum(1) # boundary / other chart + perim.scatter_add_(0, c_acc, dper) + area.scatter_add_(0, c_acc, fa[f_acc]) + nsum.index_add_(0, c_acc, fn[f_acc] * fa[f_acc, None]) + nl = nsum[:K].norm(dim=1, keepdim=True) + basis[:K] = torch.where(nl > 1e-20, nsum[:K] / nl.clamp_min(1e-20), basis[:K]) + return n_acc def segment_charts( @@ -166,162 +89,108 @@ def segment_charts( if F == 0: return torch.zeros(0, dtype=torch.long, device=device) - face_normal = mesh.face_normal.detach().cpu().numpy().astype(np.float32) - face_area = mesh.face_area.detach().cpu().numpy().astype(np.float32) - face_centroid = mesh.face_centroid.detach().cpu().numpy().astype(np.float32) - face_face = mesh.face_face.detach().cpu().numpy() + fn = mesh.face_normal.detach().to(torch.float32) + fa = mesh.face_area.detach().to(torch.float32) + fc = mesh.face_centroid.detach().to(torch.float32) + 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) - nd_cutoff = np.float32(NORMAL_DEVIATION_HARD_CUTOFF) - nd_threshold = np.float32(min(max_cost / max(w_normal_deviation, 1e-6), - NORMAL_DEVIATION_HARD_CUTOFF * 0.99)) - - component = (mesh.component.detach().cpu().numpy() - 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) + # one seed per connected component (first face of each) + comp = mesh.component.detach().long().to(device) + ncomp = int(comp.max()) + 1 if comp.numel() else 0 + if ncomp: + seeds = torch.full((ncomp,), F, dtype=torch.long, device=device) + seeds.scatter_reduce_(0, comp, torch.arange(F, device=device), reduce="amin") 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()] - if not seed_faces: - seed_faces = [0] + max_total_charts = max(F, 8000) + cap = K + F + 1 # every re-seed assigns a face, so K < K0 + F + 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) - chart_basis = np.zeros((K, 3), dtype=np.float32) - chart_normal_sum = np.zeros((K, 3), dtype=np.float32) - chart_area = np.zeros(K, dtype=np.float32) - 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()) + min_d2 = torch.full((F,), float("inf"), dtype=torch.float32, device=device) + for i in range(0, K, 32): # chunked: (F, <=32, 3) stays small + d2 = ((fc[:, None, :] - fc[seeds[i:i + 32]][None, :, :]) ** 2).sum(-1) + min_d2 = torch.minimum(min_d2, d2.amin(1)) - if K == 0: - return _renumber(face_chart, device) + # 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(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) - 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] + tq.close() # Orphan cleanup: leftover faces join their best-matching neighbor's chart. - if (face_chart == -1).any() and chart_basis.shape[0] > 0: - while True: - orphans = np.nonzero(face_chart == -1)[0] - if orphans.size == 0: - break - nbs = face_face[orphans] - 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[orphans][:, None, :] * nb_basis).sum(axis=-1) - nd = np.where(valid, np.float32(1.0) - d, np.inf) - best_e = np.argmin(nd, axis=1) - best_c = nb_charts_safe[np.arange(orphans.size), best_e] - 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, - ) + while True: + orphans = (face_chart < 0).nonzero(as_tuple=True)[0] + if orphans.numel() == 0: + break + nb = ff[orphans] + nbc = torch.where(nb >= 0, face_chart[nb.clamp_min(0)], nb.new_full((), -1)) + valid = nbc >= 0 + assignable = valid.any(1) + if not bool(assignable.any()): + break + d = (fn[orphans][:, None, :] * basis[nbc.clamp_min(0)]).sum(-1) + ndv = torch.where(valid, 1.0 - d, d.new_full((), float("inf"))) + best_c = nbc.gather(1, ndv.argmin(1, keepdim=True)).squeeze(1) + face_chart[orphans[assignable]] = best_c[assignable] + leftover = (face_chart < 0).nonzero(as_tuple=True)[0] + if leftover.numel(): # isolated faces become singleton charts + face_chart[leftover] = K + torch.arange(leftover.numel(), device=device) - return _renumber(face_chart, device) + _, inverse = torch.unique(face_chart, sorted=True, return_inverse=True) + return inverse # Parallel edge-collapse (PEC) chart clustering (GPU) @@ -400,12 +269,71 @@ def _build_chart_edges( 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( mesh: MeshData, max_cost: float = 0.7, max_iters: int = 1024, + min_faces: int = 8, ) -> 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 F = mesh.faces.shape[0] faces = mesh.faces.to(torch.long) @@ -471,5 +399,8 @@ def cluster_charts_pec( remap[win_b] = win_a 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) return inverse diff --git a/comfy_extras/nodes_mesh_postprocess.py b/comfy_extras/nodes_mesh_postprocess.py index aef7a517b..adeecb9a0 100644 --- a/comfy_extras/nodes_mesh_postprocess.py +++ b/comfy_extras/nodes_mesh_postprocess.py @@ -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 pack as _uv_pack import logging +import time from tqdm import tqdm from scipy.sparse import csr_matrix 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): """UV-unwrap a single mesh; returns (vmapping, indices, uvs); vmapping maps each output vertex to an input vertex (seam verts duplicated).""" + t_start = time.perf_counter() v_in = positions.to(torch.float32) 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) @@ -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 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() faces_np = mesh.faces.cpu().numpy() vertices_np = mesh.vertices.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_offsets_np = np.empty(n_charts + 1, dtype=np.int64) - chart_offsets_np[0] = 0 + chart_offsets_np = np.zeros(n_charts + 1, dtype=np.int64) 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 = [], [], [], [] - chart_records = [] + # progress: n_charts units for parameterize + 2*n_charts for pack (prepare + place) + 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): - gfi_np = sorted_face_idx_np[chart_offsets_np[c]:chart_offsets_np[c + 1]] - chart_faces_global = faces_np[gfi_np] - used_verts_np = np.unique(chart_faces_global) - local_faces_np = np.searchsorted(used_verts_np, chart_faces_global) - local_verts_np = vertices_np[used_verts_np] - ff_global = face_face_np[gfi_np] - ff_safe = np.maximum(ff_global, 0) - nb_chart = np.where(ff_global >= 0, face_chart_np[ff_safe], -1) - keep = (ff_global >= 0) & (nb_chart == c) - local_neighbor = np.searchsorted(gfi_np, ff_safe) - local_ff_np = np.where(keep, local_neighbor, -1) + v0, v1 = vert_offsets[c], vert_offsets[c + 1] + if ortho_ok[c]: + uvs_np_list[c] = uv0_f32[v0:v1] + continue + u = lscm_uv.get(int(c)) + if u is not None and np.all(np.isfinite(u)) and u.size: + # collapsed UV island (aspect > 100:1) blows up packing scale; keep ortho instead + bbox = u.max(axis=0) - u.min(axis=0) + if max(float(bbox.max()), 1e-12) / max(float(bbox.min()), 1e-12) <= 100.0: + uvs_np_list[c] = u + 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) - uvs = _uv_param.parametrize_chart( - torch.from_numpy(local_verts_np), lf, torch.from_numpy(local_ff_np)) - ua, ub, uc = uvs[lf[:, 0]], uvs[lf[:, 1]], uvs[lf[:, 2]] - uv_area_sum = float(0.5 * ( - (ub[:, 0] - ua[:, 0]) * (uc[:, 1] - ua[:, 1]) - - (uc[:, 0] - ua[:, 0]) * (ub[:, 1] - ua[:, 1])).abs().sum().item()) - chart_records.append({"local_faces": lf, "vmap": torch.from_numpy(used_verts_np), - "global_face_idx": torch.from_numpy(gfi_np)}) - all_chart_uvs.append(uvs) - all_chart_3d_areas.append(float(areas_cpu[c].item())) - all_chart_uv_areas.append(uv_area_sum) - all_chart_faces.append(lf) + # per-chart UV areas in one pass over all faces + uvs_all_np = np.concatenate(uvs_np_list) + ua, ub, uc = uvs_all_np[gl_faces[:, 0]], uvs_all_np[gl_faces[:, 1]], uvs_all_np[gl_faces[:, 2]] + tri_uv_area = 0.5 * np.abs( + (ub[:, 0] - ua[:, 0]) * (uc[:, 1] - ua[:, 1]) + - (uc[:, 0] - ua[:, 0]) * (ub[:, 1] - ua[:, 1])) + uv_area_np = np.bincount(chart_sorted, weights=tri_uv_area.astype(np.float64), + minlength=n_charts) + + areas_3d_np = areas_cpu.numpy().astype(np.float64) # 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 tex_per_unit = math.sqrt((target_dim * target_dim) * 0.62 / total_3d_area) - placements, atlas_w, atlas_h = _uv_pack.pack_bitmap( - all_chart_uvs, all_chart_3d_areas, all_chart_uv_areas, all_chart_faces, - texels_per_unit=tex_per_unit, padding_texels=padding) - placed = _uv_pack.apply_placements(all_chart_uvs, placements, atlas_w, atlas_h) + with tqdm(total=2 * n_charts, desc="unwrap: pack", unit="chart", leave=False) as tq_pack: + def _pack_progress(done, total): + tq_pack.update(done - tq_pack.n) + 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] out_indices = np.zeros((n_in_faces, 3), dtype=np.int64) - out_uvs_list, out_vmap_list, v_cursor = [], [], 0 - for c, rec in enumerate(chart_records): - vmap_np = rec["vmap"].cpu().numpy() - local_faces_np = rec["local_faces"].cpu().numpy() - global_face_idx = rec["global_face_idx"].cpu().numpy() - out_uvs_list.append(placed[c].cpu().numpy()) - 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) + out_indices[order] = gl_faces + vmapping_out = used_verts_all if welded_to_orig is None else welded_to_orig[used_verts_all] + uvs_out = _uv_pack.apply_placements_concat( + uvs_all_np, vert_offsets, p_x, p_y, p_sw, p_th, p_sc, p_chh, atlas_w, atlas_h) + logging.info(f"[uv_unwrap] {mesh.faces.shape[0]} faces -> {n_charts} charts, " + f"atlas {atlas_w}x{atlas_h}, {time.perf_counter() - t_start:.1f}s") return vmapping_out, out_indices, uvs_out