refactor: estimate working vae memory during encode/decode #8438
Conversation
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
psychedelicious
requested review from
lstein,
blessedcoolant,
hipsterusername and
maryhipp
as code owners
github-actions
bot
added
python
psychedelicious
force-pushed
the
psyche/refactor/estimate-vae-during-encode
branch
3 times, most recently
from
8ce9012 to
5fc38a8
Compare
hipsterusername
approved these changes
Aug 18, 2025
- Move the estimation logic to utility functions - Estimate memory _within_ the encode and decode methods, ensuring we _always_ estimate working memory when running a VAE
psychedelicious
force-pushed
the
psyche/refactor/estimate-vae-during-encode
branch
from
5fc38a8 to
9c3d4b3
Compare
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.
Summary
Extract logic to estimate working memory to utility functions. Then, instead of requiring the estimate to be passed as an arg to all encode calls (e.g. static method
ImageToLatentsInvocation.encode()), do the estimation within the encode calls. This makes it impossible to not handle the memory estimation when doing VAE stuff.Related Issues / Discussions
Also fixes an issue introduced in #8414 where some calls to VAE encode didn't pass estimate working memory: https://discord.com/channels/1020123559063990373/1149506274971631688/1405975664359903403
QA Instructions
I manually tested all code paths that do VAE stuff:
Merge Plan
n/a
Checklist
What's Newcopy (if doing a release after this PR)