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