From 73efa761d8374f688e36ede17472cb47f125dba3 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 28 Aug 2024 19:24:00 -0400 Subject: [PATCH] add taesd flux --- modules/modeldata.py | 2 ++ modules/sd_samplers.py | 2 +- modules/sd_vae_taesd.py | 31 +++++++++++++++++++++++++++---- 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/modules/modeldata.py b/modules/modeldata.py index 2ae958218..b114b26b3 100644 --- a/modules/modeldata.py +++ b/modules/modeldata.py @@ -101,6 +101,8 @@ class Shared(sys.modules[__name__].__class__): model_type = 'hunyuandit' elif "Cascade" in self.sd_model.__class__.__name__: model_type = 'sc' + elif "FluxPipeline" in self.sd_model.__class__.__name__: + model_type = 'f1' else: model_type = self.sd_model.__class__.__name__ except Exception: diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index 4db88dbe7..29034d63f 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -69,7 +69,7 @@ def create_sampler(name, model): return sampler elif shared.native: sampler = config.constructor(model) - if shared.sd_model_type == 'FluxPipeline': + if shared.sd_model_type == 'f1': if 'base_image_seq_len' not in sampler.sampler.config or 'max_image_seq_len' not in sampler.sampler.config or 'base_shift' not in sampler.sampler.config or 'max_shift' not in sampler.sampler.config: shared.log.warning('FLUX sampler: attempting to use a non compatible scheduler') return None diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index d5edbd11f..d6d41c858 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -18,6 +18,8 @@ taesd_models = { 'sdxl-encoder': None, 'sd3-decoder': None, 'sd3-encoder': None, + 'f1-decoder': None, + 'f1-encoder': None, } previous_warnings = False @@ -65,7 +67,7 @@ class TAESD2(nn.Module): # pylint: disable=abstract-method """Initialize pretrained TAESD on the given device from the given checkpoints.""" super().__init__() if latent_channels is None: - latent_channels = 16 if "taesd3" in str(encoder_path) else 4 + latent_channels = self.guess_latent_channels(str(encoder_path)) self.encoder = Encoder(latent_channels) self.decoder = Decoder(latent_channels) if encoder_path is not None: @@ -73,6 +75,14 @@ class TAESD2(nn.Module): # pylint: disable=abstract-method if decoder_path is not None: self.decoder.load_state_dict(torch.load(decoder_path, map_location="cpu")) + 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]""" @@ -92,7 +102,7 @@ class TAESD(nn.Module): # pylint: disable=abstract-method """Initialize pretrained TAESD on the given device from the given checkpoints.""" super().__init__() if latent_channels is None: - latent_channels = 16 if "taesd3" in str(encoder_path) or "taesd3" in str(decoder_path) else 4 + latent_channels = self.guess_latent_channels(str(decoder_path), str(encoder_path)) self.encoder = Encoder(latent_channels) self.decoder = Decoder(latent_channels) if encoder_path is not None: @@ -100,6 +110,14 @@ class TAESD(nn.Module): # pylint: disable=abstract-method if decoder_path is not None: self.decoder.load_state_dict(torch.load(decoder_path, map_location="cpu")) + def guess_latent_channels(self, decoder_path, encoder_path): + """guess latent channel count based on encoder filename""" + if "taef1" in encoder_path or "taef1" in decoder_path: + return 16 + if "taesd3" in encoder_path or "taesd3" in decoder_path: + return 16 + return 4 + @staticmethod def scale_latents(x): """raw latents -> [0, 1]""" @@ -148,7 +166,7 @@ def decode(latents): if model_class == 'ldm': model_class = 'sd' dtype = devices.dtype_vae if devices.dtype_vae != torch.bfloat16 else torch.float16 # taesd does not support bf16 - if 'sd' not in model_class: + if 'sd' not in model_class and 'f1' not in model_class: if not previous_warnings: previous_warnings = True shared.log.warning(f'TAESD unsupported model type: {model_class}') @@ -162,6 +180,11 @@ def decode(latents): shared.log.debug(f'VAE load: type=taesd model={model_path}') vae = taesd_models[f'{model_class}-decoder'] vae.decoder.to(devices.device, dtype) + else: + shared.log.error(f'VAE load: type=taesd model={model_path} not found') + return latents + if vae is None: + return latents try: with devices.inference_context(): latents = latents.detach().clone().to(devices.device, dtype) @@ -188,7 +211,7 @@ def encode(image): model_class = shared.sd_model_type if model_class == 'ldm': model_class = 'sd' - if 'sd' not in model_class: + if 'sd' not in model_class and 'f1' not in model_class: if not previous_warnings: previous_warnings = True shared.log.warning(f'TAESD unsupported model type: {model_class}')