xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-12 11:46:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

add taef1 handling to several places

Browse Source
This commit is contained in:
Matthew Turnshek 2024-08-16 11:53:04 -04:00
parent 05a9f3faa1
commit e792dec45f
3 changed files with 25 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
comfy/latent_formats.py
View File

@ -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

12
comfy/taesd/taesd.py
View File

@ -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]"""

14
nodes.py
View File

@ -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)