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
Open
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
…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
looks in-place but isn't -
weights + 1e-8allocates a brand-new full-size tensor before the division even starts, via the plain allocator, bypassing the disk-backed buffer above. Changed towhich 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.