Reduce UV unwrap VRAM usage

This commit is contained in:
kijai 2026-07-17 14:44:14 +03:00
parent 550af28f45
commit f40e8a099b

View File

@ -468,33 +468,39 @@ def _raster_all_torch(uvs_tex_pad, faces_pad, fmask, bw_t, bh_t, padding, device
d11 = (v1 * v1).sum(-1) d11 = (v1 * v1).sum(-1)
den = (d00 * d11 - d01 * d01).clamp(min=1e-20) den = (d00 * d11 - d01 * d01).clamp(min=1e-20)
free = comfy.model_management.get_free_memory(device)
budget = int(min(1 << 23, max(1 << 20, (free * 0.25) / 56)))
for g in sorted(set(bsz.tolist())): # one batch per pow2 grid for g in sorted(set(bsz.tolist())): # one batch per pow2 grid
sel = (bsz == g).nonzero(as_tuple=True)[0] sel_g = (bsz == g).nonzero(as_tuple=True)[0]
m = sel.shape[0] per = max(1, budget // (g * g))
xs0 = x0[sel].view(m, 1, 1) for cs in range(0, sel_g.shape[0], per):
ys0 = y0[sel].view(m, 1, 1) sel = sel_g[cs:cs + per]
cc = cid[sel] m = sel.shape[0]
bwp = bwL[cc].view(m, 1, 1) xs0 = x0[sel].view(m, 1, 1)
bhp = bhL[cc].view(m, 1, 1) ys0 = y0[sel].view(m, 1, 1)
gi = torch.arange(g, device=device) cc = cid[sel]
px = xs0 + gi.view(1, 1, g) bwp = bwL[cc].view(m, 1, 1)
py = ys0 + gi.view(1, g, 1) # (m,g,g) int bhp = bhL[cc].view(m, 1, 1)
pxf = px.float() + 0.5 gi = torch.arange(g, device=device)
pyf = py.float() + 0.5 px = xs0 + gi.view(1, 1, g)
v2x = pxf - a[sel, 0].view(m, 1, 1) py = ys0 + gi.view(1, g, 1) # (m,g,g) int
v2y = pyf - a[sel, 1].view(m, 1, 1) pxf = px.float() + 0.5
d20 = v2x * v0[sel, 0].view(m, 1, 1) + v2y * v0[sel, 1].view(m, 1, 1) pyf = py.float() + 0.5
d21 = v2x * v1[sel, 0].view(m, 1, 1) + v2y * v1[sel, 1].view(m, 1, 1) v2x = pxf - a[sel, 0].view(m, 1, 1)
idn = den[sel].view(m, 1, 1).reciprocal() v2y = pyf - a[sel, 1].view(m, 1, 1)
vv = torch.addcmul(d11[sel].view(m, 1, 1) * d20, d01[sel].view(m, 1, 1), d21, value=-1) * idn d20 = v2x * v0[sel, 0].view(m, 1, 1) + v2y * v0[sel, 1].view(m, 1, 1)
ww = torch.addcmul(d00[sel].view(m, 1, 1) * d21, d01[sel].view(m, 1, 1), d20, value=-1) * idn d21 = v2x * v1[sel, 0].view(m, 1, 1) + v2y * v1[sel, 1].view(m, 1, 1)
uu = 1.0 - vv - ww idn = den[sel].view(m, 1, 1).reciprocal()
inside = (uu >= -1e-6) & (vv >= -1e-6) & (ww >= -1e-6) vv = torch.addcmul(d11[sel].view(m, 1, 1) * d20, d01[sel].view(m, 1, 1), d21, value=-1) * idn
if padding > 0: ww = torch.addcmul(d00[sel].view(m, 1, 1) * d21, d01[sel].view(m, 1, 1), d20, value=-1) * idn
inside = _dilate_local(inside, padding) uu = 1.0 - vv - ww
valid = inside & (px < bwp) & (py < bhp) inside = (uu >= -1e-6) & (vv >= -1e-6) & (ww >= -1e-6)
flat = (cbase[cc].view(m, 1, 1) + py * bwp + px)[valid] if padding > 0:
buf[flat] = True 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 return buf, cbase