Skip to content

fix: cap tiled VAE decode output at 8GB RAM, disk-back above that, an… - #530

Open
Juan Fidel (johnnycaifan) wants to merge 1 commit into
Lightricks:masterfrom
johnnycaifan:fix/tiled-vae-decode-ram-safety
Open

fix: cap tiled VAE decode output at 8GB RAM, disk-back above that, an…#530
Juan Fidel (johnnycaifan) wants to merge 1 commit into
Lightricks:masterfrom
johnnycaifan:fix/tiled-vae-decode-ram-safety

Conversation

@johnnycaifan

Copy link
Copy Markdown

…d fix a non-in-place normalization

Both decode() and decode_spatial_temporal() preallocate their full output tensor (and decode()'s weight tensor) via torch.zeros/torch.empty sized for the entire batch (batch x frames x height x width x channels), with no cap. For a long/high-res video batch this can require tens of GB in a single allocation - on a 64GB system, running this repeatedly pushed system RAM past 97% used and required killing the process to avoid a full freeze.

Adds _allocate_decode_output(), matching the same threshold shape as the RTX Upscaler & Refiner fix in DaSiWa-Nodes: allocate normally under 8GB, back the tensor with torch.from_file() (checking real free disk space first) between 8GB and 64GB, and refuse cleanly above that instead of risking an OOM freeze.

Also fixes a separate leak in decode_spatial_temporal(): the normalization step

output /= weights + 1e-8

looks in-place but isn't - weights + 1e-8 allocates a brand-new full-size tensor before the division even starts, via the plain allocator, bypassing the disk-backed buffer above. Changed to

weights.add_(1e-8)
output /= weights

which is truly in-place. On a real render this second, unprotected allocation was still enough on its own to push RAM past the same 97% threshold even with the main output buffer disk-backed.

…d fix a non-in-place normalization

Both decode() and decode_spatial_temporal() preallocate their full
output tensor (and decode()'s weight tensor) via torch.zeros/torch.empty
sized for the entire batch (batch x frames x height x width x channels),
with no cap. For a long/high-res video batch this can require tens of
GB in a single allocation - on a 64GB system, running this repeatedly
pushed system RAM past 97% used and required killing the process to
avoid a full freeze.

Adds _allocate_decode_output(), matching the same threshold shape as
the RTX Upscaler & Refiner fix in DaSiWa-Nodes: allocate normally under
8GB, back the tensor with torch.from_file() (checking real free disk
space first) between 8GB and 64GB, and refuse cleanly above that
instead of risking an OOM freeze.

Also fixes a separate leak in decode_spatial_temporal(): the
normalization step

    output /= weights + 1e-8

looks in-place but isn't - `weights + 1e-8` allocates a brand-new
full-size tensor before the division even starts, via the plain
allocator, bypassing the disk-backed buffer above. Changed to

    weights.add_(1e-8)
    output /= weights

which is truly in-place. On a real render this second, unprotected
allocation was still enough on its own to push RAM past the same
97% threshold even with the main output buffer disk-backed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants