This repository contains the codebase for MAP-GA, which is built on the codebase of Consistency Models, and utilizes their pre-trained models and the inference pipeline.
Clone our repository and follow consistency_models for setting up the dependencies. Our modified codebase does not use Distributed PyTorch but rather launches mutliple single-GPU jobs (since it is more efficient to run several ablations this way in terms of GPU wait-time/allocation), so other dependencies for using PyTorch Distributed, such as mpi4py etc., can be ignored in our case.
We use torchmetrics based FID, LPIPS, SSIM and PSNR metrics for evaluation.
pip install torchmetricsAs a first solution, please refer to the issues of consistency_models for problems with the environment setup, especially for those regarding FlashAttention.
ImageNet diffusion models are class-conditional, so we need class information along with the images. We extract 1K images, each from a different class, from the original ImageNet validation set. This is referred to as the ImageNet1K validation set. The file imagenet_val_1k.txt contains filename and class label for 1000 images.
1. Download the original ImageNet validation dataset and extract all images into a single directory
2. Check the script create_val1k_folder.py and add the correct paths for root_dir, meta_file, dst_dir and simply run python create_val1k_folder.py
3. After running, check that dst_dir contains 1K Images each belonging to a different class.
The dst_dir path later needs to be passed as an argument if you'd like to run MAP-GA on ImageNet. (more on this later)
In the paper, we use 100 randomly picked images from the LSUNCat dataset. These are all Cat images, so the models are not class-conditional and so we just need a bunch of images in a folder. You can pick arbitrary number of images and place it in a folder whose path later needs to be passed as an argument if you'd like to run MAP-GA on LSUNCat. (more on this later)
Download pre-trained EDM models and CD models (l2 metric) for ImageNet64 and LSUNCat from consistency_models repository and put them in a folder, whose path later needs to be passed as an argument to run MAP-GA.
Note that we do not use LPIPS-trained CD models or CT models, since the former interferes with evaluation (as we use LPIPS metric for evaluation) and the latter doesn't perform well compared to CD (l2 metric) models, as reported in the consistency models paper. However, you can also just use those checkpoint paths of CT and LPIPS-trained CD models instead if you are curious.
1. Refer to the comments within scripts/job_inpaint_imagenet64c_sbatch.sh and make necessary updates
2. Refer to the comments within scripts/job_inpaint_imagenet64c_all_ablations.sh and make necessary updates
(The scripts are quite self-explanatory and have detailed comments on how to set up ablations that can run parallelly)
3. cd scripts
4. bash job_inpaint_imagenet64c_all_ablations.sh # To run MAP-GA or PiGDM or DDRM or any ablations on imagenet64
You can run several jobs for many ablations parallelly by setting up the configurations within these bash scripts. Please refer to 6363 the comments in these bash scripts for further details on setting up the configurations for MAP-GA and baselines. The bash scripts in-turn call the slurm job scripts (such as job_inpaint_imagenet64c_sbatch.sh) via sbatch. If you are running running locally or interactively on a GPU server, check the comments within these bash scripts on how to modify them accordingly.
Feel free to raise an issue to request for scripts to run MAP-GA on LSUNCat256.
After all the ablations set up in the bash scripts above have finished running, you end up with restored images in the savedir folder (default path: ./expout/imagenet64c). Then you can use the following evaluation scripts to evaluate all those ablations. The evaluation script below just spans one job and evaluates each ablation sequentially. (since evaluation is much faster than running MAP-GA or baselines)
1. Check the comments within scripts/job_evaluate_imagenet64c_all_ablations.sh
for setting up which ablations to evaluate and for any other changes
2. cd scripts
3. To run evaluation, If you are using slurm:
sbatch job_evaluate_imagenet64c_all_ablations.sh
If you are running locally or interactively on a GPU server:
bash job_evaluate_imagenet64c_all_ablations.sh >> eval_imagenet64c.out
Feel free to raise an issue to request for scripts to evaluate MAP-GA on LSUNCat256.
If you find this method and/or code useful, please consider citing
@InProceedings{Gutha_2025_WACV,
author = {Gutha, Sai Bharath Chandra and Vinuesa, Ricardo and Azizpour, Hossein},
title = {Inverse Problems with Diffusion Models: A MAP Estimation Perspective},
booktitle = {Proceedings of the Winter Conference on Applications of Computer Vision (WACV)},
month = {February},
year = {2025},
pages = {4153-4162}
}