/pytorch_diffusion_model_celebahq

PyTorch implementation of denoising diffusion probabilistic models.

Primary LanguagePythonMIT LicenseMIT

PyTorch implementation of denoising diffusion probabilistic models on the celebahq (256 * 256) dataset.

Training:

First, download the celebahq 256 * 256 dataset in this link. Rename the folder to 'CELEBAHQ' and save it under data/. Then, the images paths are going to be data/CELEBAHQ/all/00001.img ~ 30000.img

Train from scratch:

Run python main.py -c config.json -t 'train'

Train from a pretrained model:

This pretrained model has been trained for 990 epochs (~450 hours). Due to computational limit we use batchsize = 16, while the original implementaion uses batchsize = 64 . To use this checkpoint, download it (~1.59GB) and put it under model/celebahq/. Then, run python main.py -c config.json -t 'train'

Generation:

Download the pretrained model if you don't have one. Then run python main.py -c config.json -t 'generate'. Four images are generated by default. This number can be adjusted by changing gen_config["n"] in config.json.

Dependencies:

Necessary packages include Torch, Torchvision, and NumPy. A GPU is needed. If your GPU memory is less than 11GB, then you might need to decrease batchsize train_config["batch_size"] to 1 in training and decrease the number of images gen_config["n"] in generation.