Skip to content

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
invoke-ai:mainfrom
Pfannkuchensack:feat/tae-previews-video
Open

Pfannkuchensack wants to merge 15 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/tae-previews-video

Conversation

@Pfannkuchensack

@Pfannkuchensack Pfannkuchensack commented Oct 26, 2026 •

Copy link
Copy Markdown
Member

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). With tiny_autoencoder_previews: true and the matching starter installed:

Starter Base Serves
TAEW2.1 Wan, 16 ch Wan 2.2 A14B, Qwen-Image (incl. Edit), Krea-2, Anima
TAEW2.2 Wan, 48 ch Wan 2.2 TI2V-5B
TAELTX2.3 LTX-2 LTX-2.5

How it works:

  • TAEHV joins the tiny_autoencoder model 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.
  • A single latent decodes as a one-frame clip: the causal warm-up frames are trimmed and the one remaining frame is kept.
  • Wan video previews pass a 2-frame window (the middle latent frame plus the one before it). The TAE uses the earlier frame as temporal context. The linear preview still shows the same middle frame as before, pixel-identical.
  • Preview caps are per decoder: 512 px for TAEW2.1, whose decode is the heaviest, and 1024 px for all others.
  • Some TAE networks share their shapes with decoders for other models, so identification needs a name or an explicit base:
    • taeltx_2 has the same shapes as taeltx2_3 but previews LTX-2.5 worse than linear. It is refused, even with a base override.
    • taew2_1 has the same shapes as the HunyuanVideo (taehv) and CogVideoX (taecvx) decoders. Files named taehv/taecvx are refused.
    • An unnamed file with one of these networks needs the TAE's name or an explicit base. Starters always set the base.
  • When no TAE is installed, the log line now names the starter to install and the base it is listed under. For example, TAEW2.1 is listed under Wan.

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 4 over tests/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.py and tests/test_docs_json_export.py: 7377 passed.
  • ruff check and ruff format --check are clean.
  • API package: prettier and tests pass. Frontend: tsc passes.
  • Docs: pnpm test and pnpm build pass, including link validation.
  • The vendored decoder output was compared to upstream taehv.py on 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:

Family TAE TAE Linear (upscaled)
Qwen-Image / Krea-2 TAEW2.1 41.6 dB 26.6 dB
Anima / Wan image TAEW2.1 41.2 dB 26.6 dB
Wan 2.2 TI2V-5B TAEW2.2 34.4 dB 26.3 dB
LTX-2.5 TAELTX2.3 37.8 dB 26.6 dB
LTX-2.5 taeltx_2 (refused) 11.0 dB 26.6 dB
Wan video, middle frame alone / with previous frame TAEW2.1 20.3 / 34.9 dB 18.3 dB
Wan 5B video, alone / with previous frame TAEW2.2 19.3 / 27.5 dB 18.1 dB
LTX-2.5 video, middle frame TAELTX2.3 36.6 dB 27.6 dB

Cost per preview at the caps:

TAE Cap VRAM Time
TAEW2.1 512 px ~0.54 GB 10–22 ms
TAEW2.1 uncapped, 1024² (for comparison) — 2.1 GB 63 ms
TAEW2.2 1024 px ~0.54 GB ~24 ms
TAELTX2.3 1024 px ~0.27 GB ~29 ms
Wan video 2-frame window 832×480 – 1280×720 ~0.6 GB ≤ 16 ms
Wan video 2-frame window 1024² ~1.06 GB ≤ 32 ms

Browser:

  • An Anima generation previews through TAEW2.1 at 512 px, with no warnings.
  • TAEW2.1 is not offered in Anima's VAE picker.
  • Installing taeltx_2 is refused with a message pointing to TAELTX2.3.

Review

Resolved:

  • Unnamed LTX decoders and taehv/taecvx files are no longer accepted silently.
  • Window downsampling folds frames into the batch and uses 2-D pooling, so it avoids 3-D pooling kernels that may be missing on some backends.
  • The window's doubled memory is measured and documented.
  • The MiniMax H3 docstring now gives the real reasons it keeps its own path.
  • The Wan preview window and the TAELTX2.3 cap are covered by tests.
  • The starter lookup stays out of the model configs.

Known limits:

  • Qwen-Image, Krea-2 and Wan pass the noisy latent to the step callback, so their TAE previews show sharp noise in early steps. Anima and LTX-2.5 preview the predicted image.
  • For very short TI2V-5B image-to-video clips, the context frame can be the conditioned first frame, which drifts during sampling. This affects only the preview.

Compatibility / Rollout

  • No new model type and no DB migration. New config classes join the existing tiny_autoencoder type, and openapi.json and schema.ts are regenerated.
  • MiniMax H3's decoder module keeps its public names and output.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Meaningful regression coverage added / updated where needed; obsolete tests/code removed
  • Persisted-state and API changes include required migrations / compatibility validation
  • Relevant performance/efficiency opportunities considered; material claims have evidence
  • Material review findings resolved and relevant checks rerun
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

View with [code]smith Autofix with [code]smith
Need help on this PR? Tag @codesmith-bot with what you need.

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.
@github-actions github-actions Bot added python PRs that change python files invocations PRs that change invocations backend PRs that change backend files services PRs that change app services frontend PRs that change frontend files python-tests PRs that change python tests docs PRs that change docs labels Oct 26, 2026
@lstein lstein self-assigned this Oct 27, 2026
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

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

backend PRs that change backend files docs PRs that change docs frontend PRs that change frontend files invocations PRs that change invocations python PRs that change python files python-tests PRs that change python tests services PRs that change app services

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants