Repository navigation
feat(previews): Tiny AutoEncoder previews for Wan, Qwen-Image, Krea-2, Anima and LTX-2.5 - #10041
Open
Pfannkuchensack wants to merge 15 commits into
Open
Pfannkuchensack wants to merge 15 commits into
Pfannkuchensack wants to merge 15 commits into
Conversation
Opt-in `tiny_autoencoder_previews` decodes previews with an installed madebyollin TAE instead of the linear projection, via a new `tiny_autoencoder` model type and starters for SD 1.x/2.x, SDXL, SD3, FLUX.1/Z-Image, FLUX.2/ERNIE-Image and Qwen-Image-2.1. The TAE is matched by latent space and locked once per invocation; previews fall back to linear when it is missing or not resident, and frames are paced to keep their cost near a tenth of denoise time.
# Conflicts: # invokeai/frontend/api/schema.ts
…, Anima and LTX-2.5 Adds TAEHV decoders (taew2_1, taew2_2, taeltx2_3) to the `tiny_autoencoder` model type with starters, sharing one decoder base with TAESD; MiniMax H3's taeh3 decoder builds on it unchanged. Wan video previews decode the middle latent frame with the previous frame as context; preview caps are per decoder (512 px for taew2_1).
…ow preview Unnamed LTX decoders need the taeltx2_3 name or an ltx-2 base, and taehv/taecvx files are refused like taeltx_2; the missing-TAE hint names the starter. Window downsampling folds frames into the batch, and the Wan window cost and MiniMax H3 rationale are documented.
Downsample non-divisible latents on the CPU on MPS; log each fallback cause once per process. Reserve the decode's working memory with the TAE lock; SD3 and Qwen-Image-2.1 send step 0 inside the model lock. One preview object per invocation renders and paces every frame, linear when the setting is off.
# Conflicts: # invokeai/app/invocations/qwen_image_2_1/qwen_image_2_1_denoise.py # invokeai/frontend/api/schema.ts # tests/app/invocations/test_qwen_image_2_1_nodes.py
The pre-lock fit check keeps the working memory the TAE lock will reserve; make_room_in_vram takes it. The decode reserves three live activations from the decoder's structure plus two of measured conv workspace. Failures log once per TAE.
Each decoder bounds its decode: TAESD 3 live activations plus 2 of workspace, TAEHV 2.25 plus 2 per output frame. Qwen-Image and Krea-2, previewed by TAEW2.1, send step 0 inside the transformer lock.
Every decoder's bound adds a fixed 5 MiB allowance; the workspace docs name the cuDNN measurement and the untested backends. The user docs give the reserved and measured MiB for Wan video windows; Qwen-Image tests check step 0 inside both locks.
…local install names The TAE fit check offloads under the model-load read lock, and the TAE loader unwraps `state_dict` checkpoints like identification does. Re-identification of a moved local install reads its source file and folder name, so TAE, VAE and PiD records keep the base their install read.
…-previews-video # Conflicts: # invokeai/app/util/tiny_autoencoder_previews.py # tests/app/util/test_tiny_autoencoder_previews.py
This branch has not been deployed
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.
Summary
Stacked on #10040. Review and merge that PR first; this diff includes its commits until it lands.
#10040 added opt-in Tiny AutoEncoder (TAE) live previews for the image families that madebyollin's TAESD-style decoders cover. This PR extends them to the families whose VAE is the Wan or LTX video VAE, using madebyollin's TAEHV decoders (MIT, pinned to commit
62f7591). Withtiny_autoencoder_previews: trueand the matching starter installed:How it works:
tiny_autoencodermodel type through a decoder base class shared with TAESD. MiniMax H3's built-in taeh3 decoder now builds on it, with its behaviour pinned bit-for-bit by a reference test.taeltx_2has the same shapes astaeltx2_3but previews LTX-2.5 worse than linear. It is refused, even with a base override.taew2_1has the same shapes as the HunyuanVideo (taehv) and CogVideoX (taecvx) decoders. Files namedtaehv/taecvxare refused.CogView4 and Ideogram 4 keep linear previews. MiniMax H3 keeps its own decoder.
Related Issues / Discussions
Follows #10040.
QA Instructions
Automated:
uv run --no-sync pytest -n 4overtests/backend/tiny_autoencoder,tests/backend/model_manager,tests/backend/architectures,tests/backend/util,tests/model_identification,tests/backend/minimax_h3,tests/app/util,tests/app/services/progress_previews,tests/app/services/session_processor,tests/app/invocations, the preview router and socket tests,tests/test_config.pyandtests/test_docs_json_export.py: 7377 passed.ruff checkandruff format --checkare clean.tscpasses.pnpm testandpnpm buildpass, including link validation.taehv.pyon identical weights and is bit-equal.Measured on an RTX 4090, fp16. PSNR is computed against the real VAE decode of the same latent, using the tensor each family's step callback receives:
Cost per preview at the caps:
Browser:
taeltx_2is refused with a message pointing to TAELTX2.3.Review
Resolved:
taehv/taecvxfiles are no longer accepted silently.Known limits:
Compatibility / Rollout
tiny_autoencodertype, andopenapi.jsonandschema.tsare regenerated.Checklist
What's Newcopy (if doing a release after this PR)Need help on this PR? Tag
@codesmith-botwith what you need.