Skip to content

Commit cf94b6a

Browse files
committed
create a script to train vae
1 parent aeac0a0 commit cf94b6a

File tree

3 files changed

+1029
-0
lines changed

3 files changed

+1029
-0
lines changed

examples/autoencoderkl/README.md

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# AutoencoderKL training example
2+
3+
## Installing the dependencies
4+
5+
Before running the scripts, make sure to install the library's training dependencies:
6+
7+
**Important**
8+
9+
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:
10+
```bash
11+
git clone https://github.com/huggingface/diffusers
12+
cd diffusers
13+
pip install .
14+
```
15+
16+
Then cd in the example folder and run
17+
```bash
18+
pip install -r requirements.txt
19+
```
20+
21+
22+
And initialize an [🤗Accelerate](https://github.com/huggingface/accelerate/) environment with:
23+
24+
```bash
25+
accelerate config
26+
```
27+
28+
## Training on ImageNet
29+
30+
```bash
31+
accelerate launch --multi_gpu --num_processes 4 --mixed_precision bf16 train_autoencoderkl.py \
32+
--pretrained_model_name_or_path stabilityai/sdxl-vae \
33+
--num_train_epochs 100 \
34+
--gradient_accumulation_steps 2 \
35+
--learning_rate 4.5e-6 \
36+
--lr_scheduler cosine \
37+
--report_to wandb \
38+
--mixed_precision bf16 \
39+
--train_data_dir /path/to/ImageNet/train \
40+
--validation_image ./image.png \
41+
--decoder_only
42+
```
Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
accelerate>=0.16.0
2+
bitsandbytes
3+
datasets
4+
diffusers
5+
huggingface_hub
6+
lpips
7+
numpy
8+
packaging
9+
Pillow
10+
taming_transformers
11+
torch
12+
torchvision
13+
tqdm
14+
transformers
15+
wandb
16+
xformers

0 commit comments

Comments
 (0)