# Image Completion via Inference in Deep Generative Models

Code builds on that released with "Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images"

# Setup
We use Python 3.8. Install the packages in `requirements.txt` and `mpi4py`. If multi-GPU training is not required, the environemt `NO_MPI` can be set to 1 instead of installing `mpi4py` (`export NO_MPI=1`). 
```
pip install -r requirements.txt
pip install mpi4py
```
We use wandb (https://wandb.ai/) to track experiments. This requires setting the environment variable "WANDB_ENTITY" to be set to a wandb username. Alternatively, experiments can be run with the option `--unobserve`.

Data download scripts (first two from "Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images")
```
./setup_cifar10.sh
./setup_ffhq256.sh
./setup_xray.sh
```

Script to download all used pretrained unconditional models
```
./download-pretrained.sh
```

# Commands for training on CIFAR-10
```
python train.py --no_ema --ema_rate=0 --epochs_per_eval=1 --grad_accumulations=1 --hps=cifar10 --lr=0.0002 --n_batch=14 --num_epochs=50 --pretrained_load_dir=pretrained/cifar10-1/  # AIPO
python train.py --no_ema --ema_rate=0 --epochs_per_eval=1 --grad_accumulations=1 --hps=cifar10 --kls=q_r2 --lr=2e-05 --n_batch=14 --num_epochs=50 --pretrained_load_dir=pretrained/cifar10-1/ --pretrained_partial_encoder=all  # AIPO-R
```
# on FFHQ-256
```
python train.py --ema_rate=0 --no_ema --grad_accumulations=4 --hps=ffhq256 --lr=0.00015 --n_batch=1 --pretrained_load_dir=pretrained/ffhq256/  # AIPO
python train.py --no_ema --ema_rate=0 --epochs_per_eval=1 --grad_accumulations=2 --hps=ffhq256 --kls=q_r2 --lr=5e-05 --n_batch=2 --num_epochs=100 --port=None --pretrained_load_dir=pretrained/ffhq256/ --pretrained_partial_encoder=all  # AIPO-R
```
# on Chest X-ray 14
```
mpiexec -n 4 python train.py --unconditional --no_ema --ema_rate=0 --fid_samples=None --hps=xray --lr=0.00015 --n_batch=2 --skip_threshold=15000 --train_encoder_decoder=all   # train unconditional VAE - we use 4 GPUs on same node
mpiexec -n 4 python train.py --no_ema --conditioning=patches --patch_size_frac=0.25 --ema_rate=0 --epochs_per_eval=1 --fid_samples=None --hps=xray --lr=0.00015 --n_batch=2 --partial_encoder_flow_layers=0 --skip_threshold=100 --port 29500 --pretrained_load_dir=PATH/TO/UNCONDITIONAL/VAE/CHECKPOINT --skip_threshold=100  # train AIPO - we use 4 GPUs on same node
```
