From 1c14fc89db7d3dc3dd12b75f318f1e91018b6457 Mon Sep 17 00:00:00 2001 From: Matthew Turnshek Date: Fri, 16 Aug 2024 12:06:33 -0400 Subject: [PATCH] remove guess_latent_channels and add latent_channels info directly to flux model --- comfy/latent_formats.py | 1 + comfy/taesd/taesd.py | 12 +----------- 2 files changed, 2 insertions(+), 11 deletions(-) diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index a4a07b6ec..ee19faeae 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -141,6 +141,7 @@ class StableAudio1(LatentFormat): latent_channels = 64 class Flux(SD3): + latent_channels = 16 def __init__(self): self.scale_factor = 0.3611 self.shift_factor = 0.1159 diff --git a/comfy/taesd/taesd.py b/comfy/taesd/taesd.py index 184441e02..ce36f1a84 100644 --- a/comfy/taesd/taesd.py +++ b/comfy/taesd/taesd.py @@ -48,11 +48,9 @@ class TAESD(nn.Module): latent_magnitude = 3 latent_shift = 0.5 - def __init__(self, encoder_path=None, decoder_path=None, latent_channels=None): + def __init__(self, encoder_path=None, decoder_path=None, latent_channels=4): """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)) @@ -62,14 +60,6 @@ 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]"""