mirror of
https://github.com/huggingface/diffusers.git
synced 2026-01-29 07:22:12 +03:00
* create a script to train vae * update main.py * update train_autoencoderkl.py * update train_autoencoderkl.py * add a check of --pretrained_model_name_or_path and --model_config_name_or_path * remove the comment, remove diffusers in requiremnets.txt, add validation_image ote * update autoencoderkl.py * quality --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
60 lines
1.7 KiB
Markdown
60 lines
1.7 KiB
Markdown
# AutoencoderKL training example
|
|
|
|
## Installing the dependencies
|
|
|
|
Before running the scripts, make sure to install the library's training dependencies:
|
|
|
|
**Important**
|
|
|
|
To make sure you can successfully run the latest versions of the example scripts, we highly recommend **installing from source** and keeping the install up to date as we update the example scripts frequently and install some example-specific requirements. To do this, execute the following steps in a new virtual environment:
|
|
```bash
|
|
git clone https://github.com/huggingface/diffusers
|
|
cd diffusers
|
|
pip install .
|
|
```
|
|
|
|
Then cd in the example folder and run
|
|
```bash
|
|
pip install -r requirements.txt
|
|
```
|
|
|
|
|
|
And initialize an [🤗Accelerate](https://github.com/huggingface/accelerate/) environment with:
|
|
|
|
```bash
|
|
accelerate config
|
|
```
|
|
|
|
## Training on CIFAR10
|
|
|
|
Please replace the validation image with your own image.
|
|
|
|
```bash
|
|
accelerate launch train_autoencoderkl.py \
|
|
--pretrained_model_name_or_path stabilityai/sd-vae-ft-mse \
|
|
--dataset_name=cifar10 \
|
|
--image_column=img \
|
|
--validation_image images/bird.jpg images/car.jpg images/dog.jpg images/frog.jpg \
|
|
--num_train_epochs 100 \
|
|
--gradient_accumulation_steps 2 \
|
|
--learning_rate 4.5e-6 \
|
|
--lr_scheduler cosine \
|
|
--report_to wandb \
|
|
```
|
|
|
|
## Training on ImageNet
|
|
|
|
```bash
|
|
accelerate launch train_autoencoderkl.py \
|
|
--pretrained_model_name_or_path stabilityai/sd-vae-ft-mse \
|
|
--num_train_epochs 100 \
|
|
--gradient_accumulation_steps 2 \
|
|
--learning_rate 4.5e-6 \
|
|
--lr_scheduler cosine \
|
|
--report_to wandb \
|
|
--mixed_precision bf16 \
|
|
--train_data_dir /path/to/ImageNet/train \
|
|
--validation_image ./image.png \
|
|
--decoder_only
|
|
```
|