From e792dec45fe8d23b2fd90493c58288ba62d58efd Mon Sep 17 00:00:00 2001
From: Matthew Turnshek <matthewturnshek@gmail.com>
Date: Fri, 16 Aug 2024 11:53:04 -0400
Subject: [PATCH] add taef1 handling to several places

---
 comfy/latent_formats.py |  1 +
 comfy/taesd/taesd.py    | 12 +++++++++++-
 nodes.py                | 14 +++++++++++++-
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py
index ecb03b010..a4a07b6ec 100644
--- a/comfy/latent_formats.py
+++ b/comfy/latent_formats.py
@@ -162,6 +162,7 @@ class Flux(SD3):
             [-0.0005, -0.0530, -0.0020],
             [-0.1273, -0.0932, -0.0680]
         ]
+        self.taesd_decoder_name = "taef1_decoder"
 
     def process_in(self, latent):
         return (latent - self.shift_factor) * self.scale_factor
diff --git a/comfy/taesd/taesd.py b/comfy/taesd/taesd.py
index ce36f1a84..184441e02 100644
--- a/comfy/taesd/taesd.py
+++ b/comfy/taesd/taesd.py
@@ -48,9 +48,11 @@ class TAESD(nn.Module):
     latent_magnitude = 3
     latent_shift = 0.5
 
-    def __init__(self, encoder_path=None, decoder_path=None, latent_channels=4):
+    def __init__(self, encoder_path=None, decoder_path=None, latent_channels=None):
         """Initialize pretrained TAESD on the given device from the given checkpoints."""
         super().__init__()
+        if latent_channels is None:
+            latent_channels = self.guess_latent_channels(str(encoder_path))
         self.taesd_encoder = Encoder(latent_channels=latent_channels)
         self.taesd_decoder = Decoder(latent_channels=latent_channels)
         self.vae_scale = torch.nn.Parameter(torch.tensor(1.0))
@@ -60,6 +62,14 @@ class TAESD(nn.Module):
         if decoder_path is not None:
             self.taesd_decoder.load_state_dict(comfy.utils.load_torch_file(decoder_path, safe_load=True))
 
+    def guess_latent_channels(self, encoder_path):
+        """guess latent channel count based on encoder filename"""
+        if "taef1" in encoder_path:
+            return 16
+        if "taesd3" in encoder_path:
+            return 16
+        return 4
+
     @staticmethod
     def scale_latents(x):
         """raw latents -> [0, 1]"""
diff --git a/nodes.py b/nodes.py
index 16f5c9b00..2932b2288 100644
--- a/nodes.py
+++ b/nodes.py
@@ -665,6 +665,8 @@ class VAELoader:
         sd1_taesd_dec = False
         sd3_taesd_enc = False
         sd3_taesd_dec = False
+        f1_taesd_enc = False
+        f1_taesd_dec = False
 
         for v in approx_vaes:
             if v.startswith("taesd_decoder."):
@@ -679,12 +681,18 @@ class VAELoader:
                 sd3_taesd_dec = True
             elif v.startswith("taesd3_encoder."):
                 sd3_taesd_enc = True
+            elif v.startswith("taef1_encoder."):
+                f1_taesd_dec = True
+            elif v.startswith("taef1_decoder."):
+                f1_taesd_enc = True
         if sd1_taesd_dec and sd1_taesd_enc:
             vaes.append("taesd")
         if sdxl_taesd_dec and sdxl_taesd_enc:
             vaes.append("taesdxl")
         if sd3_taesd_dec and sd3_taesd_enc:
             vaes.append("taesd3")
+        if f1_taesd_dec and f1_taesd_enc:
+            vaes.append("taef1")
         return vaes
 
     @staticmethod
@@ -712,6 +720,10 @@ class VAELoader:
         elif name == "taesd3":
             sd["vae_scale"] = torch.tensor(1.5305)
             sd["vae_shift"] = torch.tensor(0.0609)
+        elif name == "taef1":
+            # TODO: check these numbers
+            sd["vae_scale"] = torch.tensor(1.5305)
+            sd["vae_shift"] = torch.tensor(0.0609)
         return sd
 
     @classmethod
@@ -724,7 +736,7 @@ class VAELoader:
 
     #TODO: scale factor?
     def load_vae(self, vae_name):
-        if vae_name in ["taesd", "taesdxl", "taesd3"]:
+        if vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]:
             sd = self.load_taesd(vae_name)
         else:
             vae_path = folder_paths.get_full_path("vae", vae_name)