This repository is our codebase for BayesDiff.
conda create --name BayesDiff python==3.8
conda activate BayesDiff
conda install pip
git clone https://github.com/karrykkk/BayesDiff.git
cd BayesDiff
pip install -r requirements.txt
This repository integrates uncertainty quantification into three models, each in its own folder:
Each folder contains a custom_model.py
that emerged with uncertainty quantification techniques.
cd ddpm_and_guided
your_local_model_path
your_local_model_path
.
cd configs
vim imagenet128_guided.yml
your_local_image_path
.self.image_path
attribute of class imagenet_dataset in la_train_datasets.py
to your_local_image_path
.
vim la_train_datasets.py
In the file dpm.sh
, you will find a template for usage for UQ-itegrated dpm-solver-2 sampler. By running this bash script, you can get the sorted_sample.png
based on the image-wise uncertainty metric.
bash dpm.sh
For other samplers, just change dpm.sh
to ddpm.sh
or ddim.sh
.
cd sd
Download Stable Diffusion v1.5 to your_local_model_path
Please download subset of laion-art to your_local_image_path
. These images is a subset from the LAION-Art dataset, store it in your_laion_art_path
. This will allow you to retrieve the corresponding prompts for the downloaded images. Note that a subset of approximately 1000 images is sufficient for effectively fitting the LLLA.
In the file sd.sh
, you will find a template for usage. Please adjust this template to match your local file path and the specific prompt you intend to use.
bash sd.sh
cd uvit
your_local_encoder_path
. Download ImageNet 256x256 (U-ViT-H/2) to your_local_uvit_path
.
your_local_image_path
.self.image_path
attribute of class imagenet_feature_dataset in la_train_datasets.py
to your_local_image_path
.
vim la_train_datasets.py
In the file dpm.sh
, you will find a template for usage. Please adjust this template to match your local file path.
bash dpm.sh
This codebase is based on remarkable projects from the community, including DPM-Solver, U-ViT, Stable Diffusion.
If you find out work useful, please cite our paper at:
@inproceedings{kou2023bayesdiff,
title={BayesDiff: Estimating Pixel-wise Uncertainty in Diffusion via Bayesian Inference},
author={Kou, Siqi and Gan, Lei and Wang, Dequan and Li, Chongxuan and Deng, Zhijie},
booktitle={The Twelfth International Conference on Learning Representations},
year={2023}
}