Diffusion Autoencoders

This is an unofficial implementation of Diffusion Autoencoders. You can find the official implementation here.

🛠️ Setup

If you are using poetry, you can install the required packages by running the following command:

poetry install

For more details on the required packages, refer to pyproject.toml.

🚀 Run

Training Diff-AE

You can train Diff-AE on CelebA or CelebA-HQ datasets.
There is no need to download the datasets manually; they are automatically downloaded when you run.

export DATA_NAME="celebahq"
export IMAGE_SIZE=128
export EXPN="hoge"

poetry run diffae_train \
    --data_name=$DATA_NAME \
    --image_size=$IMAGE_SIZE \
    --expn=$EXPN

You can modify the settings before training. If you wish to create custom settings, modify the diffae/cfg/{IMAGE_SIZE}_model.yml file according to your GPU specs. ⚙️

The training results will be saved in output/{EXPN}. If the --expn argument is not provided, the directory name will be generated based on the current time.

Note that Latent DDIM is not implemented as unconditional image synthesis by Diff-AE is not the focus of this repo.

Evaluating Diff-AE

You can evaluate the trained model by calculating its MSE and LPIPS.

export PATH_TO_OUTPUT_DIR="output/hoge"
export MODEL_CKPT="last_ckpt.pth"

poetry run diffae_test \
    --output=$PATH_TO_OUTPUT_DIR \
    --model_ckpt=$MODEL_CKPT

Training the Classifier

You can also train a classifier for attribute manipulation.

export PATH_TO_OUTPUT_DIR="output/hoge"
export MODEL_CKPT="last_ckpt.pth"

poetry run clf_train \
    --output=$PATH_TO_OUTPUT_DIR \
    --model_ckpt=$MODEL_CKPT

Evaluating the Classifier

The classifier can be evaluated by calculating its accuracy and AUROC.

export PATH_TO_OUTPUT_DIR="output/hoge"
export MODEL_CKPT="last_ckpt.pth"
export CLF_CKPT="clf_last_ckpt.pth"

poetry run clf_test \
    --output=$PATH_TO_OUTPUT_DIR \
    --model_ckpt=$MODEL_CKPT \
    --clf_ckpt=$CLF_CKPT

📚 Example Notebooks

You can check out the minimal working examples in the notebooks found here.

If you find this repository helpful, please consider giving a star ⭐!

otodd07 / diffusion-autoencoders Goto Github PK