Skip to content

fix(model_manager): detect Flux1/2 VAE by latent space dimensions instead of filename#8790

Merged
lstein merged 6 commits intoinvoke-ai:mainfrom
Pfannkuchensack:fix/vae-flux-detection-probing
Jan 27, 2026
Merged

fix(model_manager): detect Flux1/2 VAE by latent space dimensions instead of filename#8790
lstein merged 6 commits intoinvoke-ai:mainfrom
Pfannkuchensack:fix/vae-flux-detection-probing

Conversation

@Pfannkuchensack
Copy link
Copy Markdown
Collaborator

@Pfannkuchensack Pfannkuchensack commented Jan 27, 2026

Summary

VAE detection previously relied on filename pattern matching and BatchNorm layer detection, which could fail for VAE files with generic names. Now probes the model's decoder.conv_in weight shape to determine the latent space dimensions:

  • 32 channels -> Flux2 VAE
  • 16 channels -> Flux1 VAE
  • 4 channels -> SD/SDXL VAE (with filename fallback for SD1/SD2/SDXL distinction)

Flux2 VAE detection now uses two identification methods:

  1. BatchNorm layers (bn.running_mean, bn.running_var) - unique to Flux2
  2. 32-dimensional latent space (decoder.conv_in with 32 input channels)

Related Issues / Discussions

N/A

QA Instructions

  1. Get a Flux1 VAE safetensors file with a generic name (e.g. ae.safetensors)
  2. Get a Flux2 VAE safetensors file
  3. Try to install/import both VAEs via the model manager
  4. Verify Flux1 VAE is correctly detected as Flux base type
  5. Verify Flux2 VAE is correctly detected as Flux2 base type

Merge Plan

N/A - straightforward change, can merge anytime.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Tests added / updated (if applicable)
  • ❗Changes to a redux slice have a corresponding migration
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

…d of filename

VAE detection previously relied solely on filename pattern matching, which failed
for Flux VAE files with generic names like "ae.safetensors". Now probes the model's
decoder.conv_in weight shape to determine the latent space dimensions:
- 16 channels -> Flux VAE
- 4 channels -> SD/SDXL VAE (with filename fallback for SD1/SD2/SDXL distinction)
@github-actions github-actions bot added python PRs that change python files backend PRs that change backend files labels Jan 27, 2026
Extend Flux2 VAE detection to also check for 32-dimensional latent space
(decoder.conv_in with 32 input channels) in addition to BatchNorm layers.
This provides more robust detection for Flux2 VAE files regardless of filename.
@Pfannkuchensack Pfannkuchensack changed the title fix(model_manager): detect Flux VAE by latent space dimensions instead of filename fix(model_manager): detect Flux1/2 VAE by latent space dimensions instead of filename Jan 27, 2026
Copy link
Copy Markdown
Collaborator

@lstein lstein left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tested on VAEs from flux.1, flux.2, SDXL and ZiT. Works as advertised.

@lstein lstein enabled auto-merge (squash) January 27, 2026 05:15
@lstein lstein merged commit 0ae7392 into invoke-ai:main Jan 27, 2026
13 checks passed
@Pfannkuchensack Pfannkuchensack deleted the fix/vae-flux-detection-probing branch January 27, 2026 06:00
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 python PRs that change python files

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants