mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-18 20:38:15 +08:00
Reduce UV unwrap VRAM usage
This commit is contained in:
parent
550af28f45
commit
f40e8a099b
@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user