From 8dc19e40d129c8ee049be7be2657458509717ba5 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 24 Apr 2024 09:20:31 -0400 Subject: [PATCH] Don't init a VAE model when there are no VAE weights. --- comfy/sd.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/comfy/sd.py b/comfy/sd.py index 57dba0b44..16dc0b732 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -210,7 +210,7 @@ class VAE: self.first_stage_model = StageC_coder() self.downscale_ratio = 32 self.latent_channels = 16 - else: + elif "decoder.conv_in.weight" in sd: #default SD1.x/SD2.x VAE parameters ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0} @@ -226,6 +226,10 @@ class VAE: self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"}, encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': ddconfig}, decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': ddconfig}) + else: + logging.warning("WARNING: No VAE weights detected, VAE not initalized.") + self.first_stage_model = None + return else: self.first_stage_model = AutoencoderKL(**(config['params'])) self.first_stage_model = self.first_stage_model.eval()