Leminhbinh0209 / FinetuneVAE-SD

Fine-tune VAE of Stable Diffusion model
8 stars 1 forks source link

This repository is created to fine-tune your VAE model of Stable Diffusion model, which you can change input image size, or with a new dataset.

Note: I follow the guidance here, in which some first epochs are trained with (l1 + Lpips), later epochs are trained with (l2 + 0.1*Lpips) loss.

Please download pre-trained SD model here, and put into sd_model folder.

Copy ldm folder from this repo into your current directory.

For training, run the following script:

python train.py \
    --data_dir <YOUR DATA FOLDER>
    --batch_size 2 \
    --num_epochs 20 \
    --lr 2e-5 \
    --val_size 0.1 \
    --precision 16 \
    --image_size <YOUR DESIRED SIZE> \
    --lpips_loss_weight 1.0 \
    --ema_decay 0.99\