feat(mm): improved VAE encode VRAM usage - #8414
Merged
Merged
Conversation
psychedelicious
requested review from
blessedcoolant,
hipsterusername,
lstein and
maryhipp
as code owners
August 11, 2025 09:07
1 task
This commit includes a task delegated to Claude to investigate our VAE working memory calculations and investigation results. See VAE_INVESTIGATION.md for motivation and detail. Everything else is its output. Result data includes empirical measurements for all supported model architectures at a variety of resolutions and fp16/fp32 precision. Testing conducted on a 4090. The summarized conclusion is that our working memory estimations for decoding are spot-on, but decoding also needs some extra working memory. Empirical measurements suggest ~45% the amount needed for encoding. A followup commit will implement working memory estimations for VAE encoding with the goal of preventing unexpected OOMs during encode.
…models Tell the model manager that we need some extra working memory for VAE encoding operations to prevent OOMs. See previous commit for investigation and determination of the magic numbers used. This safety measure is especially relevant now that we have FLUX Kontext and may be encoding rather large ref images. Without the working memory estimation we can OOM as we prepare for denoising. See #8405 for an example of this issue on a very low VRAM system. It's possible we can have the same issue on any GPU, though - just a matter of hitting the right combination of models loaded.
This reverts commit bc9ed57.
If the transformer fills up VRAM, then when we VAE encode kontext latents, we'll need to first offload the transformer (partially, if partial loading is enabled). No need to do this - we can encode kontext latents before loading the transformer to reduce model thrashing.
… in latent space We weren't tracking the canvas dimensions properly which coudl result in FLUX not "seeing" ref images after the first very well
Prevents a large spike in VRAM when preparing to denoise w/ multiple ref images. There doesn't appear to be any different in image quality / ref adherence when concatenating in latent space vs image space, though images _are_ different.
… extension Slightly reduces VRAM allocations.
psychedelicious
force-pushed
the
psyche/feat/vae-encode-working-memory
branch
from
August 11, 2025 12:56
788f7ba to
51d004e
Compare
hipsterusername
approved these changes
Aug 11, 2025
4 tasks
kappacommit
pushed a commit
to kappacommit/InvokeAI
that referenced
this pull request
Jul 8, 2026
…code Flux2RefImageExtension encoded reference images with model_on_device() and no working-memory estimate, unlike FLUX.1's KontextExtension which passes one. A single reference can be up to ~4.1MP after the cap, whose encode needs far more than the default reserve - the encode then allocates past what the cache holds free and, on Windows, silently spills to shared memory. Pre-existing on main (independent of smart_partial_loading, since op-provided estimates are the existing VAE-op mechanism); surfaced by review of this PR. Mirrors the KontextExtension estimate (encode uses ~50% of the decode scaling constant, see invoke-ai#8414). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
2 participants
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
This PR includes a reverted commit with the investigation that determined the magic numbers required for VAE encode working memory estimation. Claude did the testing - reverted commit includes prompt, test scripts, results. Our existing magic numbers used for VAE decoding were confirmed to be accurate.
Related Issues / Discussions
QA Instructions
This change uses existing, proven APIs. I don't think it will cause any issues.
Merge Plan
n/a
Checklist
What's Newcopy (if doing a release after this PR)