mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-07 23:21:04 +08:00
862 lines
37 KiB
Python
862 lines
37 KiB
Python
"""Atlas packing via bitmap rasterize-and-place."""
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
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, 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
|
|
offset: Tuple[float, float] # in texels
|
|
scale: float # texels per UV unit
|
|
rotation: float = 0.0 # radians
|
|
swap_xy: bool = False # extra 90° bitmap rotation chosen at place time
|
|
chart_h: float = 0.0 # unswapped bitmap height in texels (rotation pivot)
|
|
|
|
|
|
@_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 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(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:
|
|
xmax = xr
|
|
if yr < ymin:
|
|
ymin = yr
|
|
if yr > ymax:
|
|
ymax = yr
|
|
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
|
|
|
|
|
|
@_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 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
|
|
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
|
|
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 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 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, :]
|
|
y[:, :-1, :] |= x[:, 1:, :]
|
|
y[:, :, 1:] |= x[:, :, :-1]
|
|
y[:, :, :-1] |= x[:, :, 1:]
|
|
x = y
|
|
return x
|
|
|
|
|
|
def _raster_all_torch(uvs_tex_pad, faces_pad, fmask, bw_t, bh_t, padding, device):
|
|
"""Rasterize every chart into one flat bool buffer; buf[cbase[i]:cbase[i+1]].view(bh,bw)
|
|
is chart i's bitmap. Triangles are bucketed by next-pow2 bbox size to bound memory."""
|
|
n = uvs_tex_pad.shape[0]
|
|
fmax = faces_pad.shape[1]
|
|
bwL, bhL = bw_t.long(), bh_t.long()
|
|
cbase = torch.zeros(n + 1, dtype=torch.long, device=device)
|
|
torch.cumsum(bwL * bhL, 0, out=cbase[1:])
|
|
buf = torch.zeros(int(cbase[-1].item()), dtype=torch.bool, device=device)
|
|
|
|
# gather all triangle coords, keep only valid faces -> (Ttot,3,2) + chart id per triangle
|
|
fp = faces_pad.reshape(n, fmax * 3)
|
|
tri = torch.gather(uvs_tex_pad, 1, fp[..., None].expand(-1, -1, 2)).reshape(n * fmax, 3, 2)
|
|
fm = fmask.reshape(-1)
|
|
tri_f = tri[fm]
|
|
if tri_f.shape[0] == 0:
|
|
return buf, cbase
|
|
cid = torch.arange(n, device=device).repeat_interleave(fmax)[fm]
|
|
|
|
# per-triangle pixel bbox, inflated by padding (origin >= 0); bucket by next-pow2 max-dim
|
|
tmin = tri_f.amin(1)
|
|
tmax = tri_f.amax(1)
|
|
x0 = (tmin[:, 0].floor().long() - padding).clamp_min(0)
|
|
y0 = (tmin[:, 1].floor().long() - padding).clamp_min(0)
|
|
bbw = (tmax[:, 0].ceil().long() + padding) - x0 + 1
|
|
bbh = (tmax[:, 1].ceil().long() + padding) - y0 + 1
|
|
mxd = torch.maximum(bbw, bbh).clamp_min(1)
|
|
bsz = (2 ** torch.ceil(torch.log2(mxd.float())).long()).long()
|
|
|
|
a = tri_f[:, 0]
|
|
b = tri_f[:, 1]
|
|
c = tri_f[:, 2]
|
|
v0 = b - a
|
|
v1 = c - a
|
|
d00 = (v0 * v0).sum(-1)
|
|
d01 = (v0 * v1).sum(-1)
|
|
d11 = (v1 * v1).sum(-1)
|
|
den = (d00 * d11 - d01 * d01).clamp(min=1e-20)
|
|
|
|
for g in sorted(set(bsz.tolist())): # one batch per pow2 grid
|
|
sel = (bsz == g).nonzero(as_tuple=True)[0]
|
|
m = sel.shape[0]
|
|
xs0 = x0[sel].view(m, 1, 1)
|
|
ys0 = y0[sel].view(m, 1, 1)
|
|
cc = cid[sel]
|
|
bwp = bwL[cc].view(m, 1, 1)
|
|
bhp = bhL[cc].view(m, 1, 1)
|
|
gi = torch.arange(g, device=device)
|
|
px = xs0 + gi.view(1, 1, g)
|
|
py = ys0 + gi.view(1, g, 1) # (m,g,g) int
|
|
pxf = px.float() + 0.5
|
|
pyf = py.float() + 0.5
|
|
v2x = pxf - a[sel, 0].view(m, 1, 1)
|
|
v2y = pyf - a[sel, 1].view(m, 1, 1)
|
|
d20 = v2x * v0[sel, 0].view(m, 1, 1) + v2y * v0[sel, 1].view(m, 1, 1)
|
|
d21 = v2x * v1[sel, 0].view(m, 1, 1) + v2y * v1[sel, 1].view(m, 1, 1)
|
|
idn = den[sel].view(m, 1, 1).reciprocal()
|
|
vv = torch.addcmul(d11[sel].view(m, 1, 1) * d20, d01[sel].view(m, 1, 1), d21, value=-1) * idn
|
|
ww = torch.addcmul(d00[sel].view(m, 1, 1) * d21, d01[sel].view(m, 1, 1), d20, value=-1) * idn
|
|
uu = 1.0 - vv - ww
|
|
inside = (uu >= -1e-6) & (vv >= -1e-6) & (ww >= -1e-6)
|
|
if padding > 0:
|
|
inside = _dilate_local(inside, padding)
|
|
valid = inside & (px < bwp) & (py < bhp)
|
|
flat = (cbase[cc].view(m, 1, 1) + py * bwp + px)[valid]
|
|
buf[flat] = True
|
|
return buf, cbase
|
|
|
|
|
|
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 _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)]
|
|
|
|
# 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)
|
|
|
|
# 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)])
|
|
|
|
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, 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:
|
|
return [], 1, 1
|
|
device = comfy.model_management.get_torch_device()
|
|
ang = torch.linspace(0.0, math.pi / 2.0, 37, device=device)[:-1]
|
|
cos_a, sin_a = ang.cos(), ang.sin()
|
|
|
|
# ---- Prepare pass 1: best-rotation + scale + bbox for ALL charts at once (batched) ----
|
|
vcount = [int(u.shape[0]) for u in chart_uvs]
|
|
fcount = [int(f.shape[0]) for f in chart_faces]
|
|
vmax = max(vcount)
|
|
fmax = max(fcount)
|
|
uvs_pad = torch.zeros(n, vmax, 2, device=device)
|
|
vmask = torch.zeros(n, vmax, dtype=torch.bool, device=device)
|
|
faces_pad = torch.zeros(n, fmax, 3, dtype=torch.long, device=device)
|
|
fmask = torch.zeros(n, fmax, dtype=torch.bool, device=device)
|
|
for i in range(n):
|
|
uvs_pad[i, :vcount[i]] = chart_uvs[i].to(device=device, dtype=torch.float32)
|
|
vmask[i, :vcount[i]] = True
|
|
if fcount[i]:
|
|
faces_pad[i, :fcount[i]] = chart_faces[i].to(device=device, dtype=torch.long)
|
|
fmask[i, :fcount[i]] = True
|
|
u0, u1 = uvs_pad[..., 0], uvs_pad[..., 1] # (N,Vmax)
|
|
BIG = 1e30
|
|
mlo = torch.where(vmask, torch.zeros_like(u0), u0.new_full((), BIG))
|
|
mhi = torch.where(vmask, torch.zeros_like(u0), u0.new_full((), -BIG))
|
|
xr = torch.addcmul(u0[:, :, None] * cos_a, u1[:, :, None], sin_a, value=-1) # (N,Vmax,A)
|
|
yr = torch.addcmul(u0[:, :, None] * sin_a, u1[:, :, None], cos_a)
|
|
xsp = (xr + mhi[:, :, None]).amax(1) - (xr + mlo[:, :, None]).amin(1) # (N,A) masked span
|
|
ysp = (yr + mhi[:, :, None]).amax(1) - (yr + mlo[:, :, None]).amin(1)
|
|
ti = (xsp * ysp).argmin(1) # (N,) best angle per chart
|
|
cc, ss = cos_a[ti][:, None], sin_a[ti][:, None] # (N,1)
|
|
rx = torch.addcmul(u0 * cc, u1, ss, value=-1) # (N,Vmax)
|
|
ry = torch.addcmul(u0 * ss, u1, cc)
|
|
rxmin = (rx + mlo).amin(1) # (N,)
|
|
rxmax = (rx + mhi).amax(1)
|
|
rymin = (ry + mlo).amin(1)
|
|
rymax = (ry + mhi).amax(1)
|
|
a3 = torch.tensor([max(a, 1e-12) for a in chart_3d_areas], device=device)
|
|
au = torch.tensor([max(a, 1e-12) for a in chart_uv_areas], device=device)
|
|
base = (a3 / au).sqrt() * texels_per_unit
|
|
maxb = (4.0 * a3.sqrt() * texels_per_unit).clamp_min(8.0)
|
|
bbm = torch.maximum(rxmax - rxmin, rymax - rymin).clamp_min(1e-12)
|
|
scale = torch.minimum(base, maxb / bbm) # (N,)
|
|
uvs_tex_pad = torch.stack([(rx - rxmin[:, None]) * scale[:, None],
|
|
(ry - rymin[:, None]) * scale[:, None]], dim=-1) # (N,Vmax,2)
|
|
bw_t = ((rxmax - rxmin) * scale).ceil().int() + padding_texels + 1
|
|
bh_t = ((rymax - rymin) * scale).ceil().int() + padding_texels + 1
|
|
|
|
# one sync: pull all per-chart scalars
|
|
thetas = ang[ti].cpu().tolist()
|
|
scales = scale.cpu().tolist()
|
|
|
|
# ---- 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)
|
|
|
|
# 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)]
|
|
|
|
# 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
|
|
side_guess = int(math.sqrt(sum(d[0] * d[1] for d in dim_l)) * 2) + 16
|
|
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(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 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)
|
|
na[:atlas.shape[0], :atlas.shape[1]] = atlas
|
|
atlas = na
|
|
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)
|
|
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]
|
|
bh_, bw_ = (dimr if swap else dim)
|
|
atlas[by + pix[:, 0], bx + pix[:, 1]] = True # sparse blit
|
|
cur_w = max(cur_w, bx + bw_)
|
|
cur_h = max(cur_h, by + bh_)
|
|
ct = (ctr_pad if swap else ct_pad)[ci, :bw_] # GPU skyline lift
|
|
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(dim_l[ci][0]))
|
|
return placements, cur_w, cur_h
|
|
|
|
|
|
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,
|
|
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 empty, empty, empty, empty.astype(np.float64), empty.astype(np.float64), empty, 1, 1
|
|
if not _HAVE_NUMBA_PACK:
|
|
chart_uvs = [torch.from_numpy(np.ascontiguousarray(uvs_cat[uv_offsets[c]:uv_offsets[c + 1]]))
|
|
for c in range(n)]
|
|
chart_faces = [torch.from_numpy(np.ascontiguousarray(faces_cat[face_offsets[c]:face_offsets[c + 1]]))
|
|
for c in range(n)]
|
|
placements, w, h = _pack_bitmap_torch(
|
|
chart_uvs, [float(a) for a in chart_3d_areas], [float(a) for a in chart_uv_areas],
|
|
chart_faces, texels_per_unit, padding_texels, attempts=attempts,
|
|
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)
|
|
# 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_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))
|
|
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)
|