A PyTorch reimplementation of Denoising Diffusion Probabilistic Models with checkpoints converted from the author's TensorFlow implementation.
Running
pip install -e git+https://github.com/pesser/pytorch_diffusion.git#egg=pytorch_diffusion
pytorch_diffusion_demo
will start a Streamlit demo. It is recommended to run the demo with a GPU available.
Diffusion models with pretrained weights for cifar10
, lsun-bedroom
,
lsun_cat
or lsun_church
can be loaded as follows:
from pytorch_diffusion import Diffusion
diffusion = Diffusion.from_pretrained("lsun_church")
samples = diffusion.denoise(4)
diffusion.save(samples, "lsun_church_sample_{:02}.png")
The U-Net model used for denoising is available via diffusion.model
and can
also be instantiated on its own:
from pytorch_diffusion import Model
model = Model(resolution=32,
in_channels=3,
out_ch=3,
ch=128,
ch_mult=(1,2,2,2),
num_res_blocks=2,
attn_resolutions=(16,),
dropout=0.1)
This configuration example corresponds to the model used on CIFAR-10.
If you installed directly from github, you can find the cloned
repository
in <venv path>/src/pytorch_diffusion
for virtual environments, and
<cwd>/src/pytorch_diffusion
for global installs. There, you can run
python pytorch_diffusion/diffusion.py <name> <bs> <nb>
where <name>
is one of cifar10
, lsun-bedroom
, lsun_cat
, lsun_church
,
<bs>
is the batch size and <nb>
the number of batches. This will produce
samples from the PyTorch models and save them to results/<name>/
.
The converted pytorch checkpoints are provided for download. If you want to convert them on your own, you can follow the steps described here.
This section assumes your working directory is the root of this repository. Download the pretrained TensorFlow checkpoints. It should follow the original structure,
diffusion_models_release/
diffusion_cifar10_model/
model.ckpt-790000.data-00000-of-00001
model.ckpt-790000.index
model.ckpt-790000.meta
diffusion_lsun_bedroom_model/
...
...
Set the environment variable TFROOT
to the directory where you want to store
the author's repository, e.g.
export TFROOT=".."
Clone the diffusion repository,
git clone https://github.com/hojonathanho/diffusion.git ${TFROOT}/diffusion
and install their required dependencies
(pip install ${TFROOT}/requirements.txt
). Then add the following to your
PYTHONPATH
:
export PYTHONPATH=".:./scripts:${TFROOT}/diffusion:${TFROOT}/diffusion/scripts:${PYTHONPATH}"
To test the pytorch implementations of the required operations against their TensorFlow counterparts under random initialization and random inputs, run
python -c "import convert as m; m.test_ops()"
To load the pretrained TensorFlow models, copy the weights into the pytorch models, check for equality on random inputs and finally save the corresponding pytorch checkpoints, run
python -c "import convert as m; m.transplant_cifar10()"
python -c "import convert as m; m.transplant_lsun_bedroom()"
python -c "import convert as m; m.transplant_lsun_cat()"
python -c "import convert as m; m.transplant_lsun_church()"
Pytorch checkpoints will be saved in
diffusion_models_converted/
diffusion_cifar10_model/
model-790000.ckpt
diffusion_lsun_bedroom_model/
model-2388000.ckpt
diffusion_lsun_cat_model/
model-1761000.ckpt
diffusion_lsun_church_model/
model-4432000.ckpt
To produce N
samples from each of the pretrained TensorFlow models, run
python -c "import convert as m; m.sample_tf(N)"
Samples will be saved in results/
.