From 0e4ee04488159b81d95a9ffcd983a077fd5dcb77 Mon Sep 17 00:00:00 2001
From: Wagner Bruna <wbruna@users.noreply.github.com>
Date: Wed, 27 May 2026 22:13:16 -0300
Subject: [PATCH] fix: correct tae for models that use the flux2 vae (#1571)

---
 src/tae.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/tae.hpp b/src/tae.hpp
index 62f14c9e..f43ef48f 100644
--- a/src/tae.hpp
+++ b/src/tae.hpp
@@ -584,7 +584,7 @@ public:
     TAESD(bool decode_only = true, SDVersion version = VERSION_SD1)
         : decode_only(decode_only) {
         bool use_midblock_gn = false;
-        taef2                = sd_version_is_flux2(version);
+        taef2                = sd_version_uses_flux2_vae(version);
 
         if (sd_version_is_dit(version)) {
             z_channels = 16;