Autoencode training

[darshats]

Hi, I'm posting a new issue since other one is closed.

I'm trying this exact scenario of retraining the sdxl vae. I've made the changes indicated to create a config file that looks like below so that its pointing to training AutoencoderKL class with the loss config of lpips, and weights checkpoint of sdxlvae

The problem I run into is that the encode methods on the AutoencoderKL and its parent AutoencoderEngine are different. I'm using the training script under sgm examples. When training runs, as part of sanity check, the encode method gets called via:

run_sanity_check-->validation_step-->AutoencodingEngine.forward()-->self.encode(x, return_reg_log=True)

At this point, derived class AutoencoderKL does not have an encode method that takes return_reg_log param. What is the right codepath here to make this match sdxl vae training?

config file

model:
  base_learning_rate: 4.5e-6
  target: sgm.models.autoencoder.AutoencoderKL
  params:
    input_key: jpg
    monitor: val/rec_loss
    embed_dim: 4 
    ckpt_path: "/mnt/d/generative-models/configs/checkpoints/sdxl_vae.safetensors"

    lossconfig:
      target: sgm.modules.autoencoding.losses.GeneralLPIPSWithDiscriminator
      params:
        perceptual_weight: 0.25
        disc_start: 20001
        disc_weight: 0.5
        learn_logvar: True

        regularization_weights:
          kl_loss: 1.0

    regularizer_config:
      target: sgm.modules.autoencoding.regularizers.DiagonalGaussianRegularizer

    ddconfig: 
       attn_type: vanilla-xformers 
       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 

Thank you!

Originally posted by @darshats in https://github.com/Stability-AI/generative-models/issues/85#issuecomment-1703762148

[YAOYI626]

hey @darshats did you get any updated about this? Is there any correct config that can be used to further finetune the SDXL VAE model? Thanks in advance.

[wtliao]

@YAOYI626 Hi, have you made it? Thanks