Blue noise for diffusion models

This repository contains the code for our paper:

Blue noise for diffusion models | [project page]

Xingchang Huang, Corentin Salaün, Cristina Vasconcelos, Christian Theobalt, Cengiz Öztireli, Gurprit Singh

SIGGRAPH (Conference Proceedings), 2024

TL;DR we present:

• a first framework studying the impact of noise correlation within and across images in diffusion models;
• a novel time-varying noise model (from Gaussian white to blue noise) for denoising diffusion models;
• a fast blue noise generation method using matmul and tiling;
• evaluations on image generation, super-resolution and latent diffusion.

📰 News

• 06/2024: Code released.

📂 Code structure

• asset: includes images for the README.
• bluenoise: includes get_noise function and the lower-triangular matrix for generating Gaussian blue/red noise. Download the two matrices (.npz files) from here and put them into the bluenoise folder.
• scripts: includes scripts for reproducing the results in the paper.
• data: includes datasets for training from scratch and super-resolution test. Download the folder from here. cat_res64 and cat_res512 are for training. church_res128_test is for testing super-resolution. Other folders are dummy with few images inside, you can prepare data in the same format if you need to train on your own data.
• results_gaussianBN, results_gaussianBN_superres: includes folders of pre-trained models and intial noise for reproducing the figures in the paper. Download the folders from here. You can also only download the folders you need and find the corresponding scripts in scripts/sampling following the Usage section below.

🛠️ Installation

The following are tested on Windows 10 / Linux, with an NVIDIA GeForce RTX 4090 and CUDA 11.8 installed.

First, install Anaconda here if you haven't. Then, install the required packages in the Anaconda command line environment by running the following commands:

conda create -n bndm python=3.10
conda install pytorch==2.0.0 torchvision==0.15.0 torchaudio==2.0.0 pytorch-cuda=11.8 -c pytorch -c nvidia
pip install diffusers
pip install accelerate
pip install matplotlib
pip install scipy
pip install opencv-python
pip install lmdb

🗝️ Usage

🤗 Noise generation

Reproduce main paper figure 3, 4, inset, 10 and supplemental figure 1, 2. A folder named scripts/results will be automatically created and the results are shown there.

python scripts/fig_main_3_4_inset_10_supp_1_2.py

🤗 Sampling using pre-trained models

To reproduce main paper figure 1, 6, 9, 11, 12, we provide the corresponding scripts in scripts/sampling:

• iadb_church_superres_test.sh: figure 6, corresponds to folders in results_gaussianBN_superres
• church_res64_test.sh: figure 1, 11, corresponds to folders results_gaussianBN/church_res64_*
• cat_res64_test.sh: figure 11, corresponds to folders results_gaussianBN/cat_res64_*
• celeba_res64_test.sh: figure 11, corresponds to folders results_gaussianBN/celeba_res64_*
• cat_res128_test.sh: figure 12, corresponds to folders results_gaussianBN/cat_res128_*
• celeba_res128_test.sh: figure 12, corresponds to folders results_gaussianBN/celeba_res128_*
• latent_iadb_cat_res512_test.sh: figure 9, corresponds to folders results_gaussianBN/latent_iadb_cat_res512_*

Results are shown in images and seqs within each folder.

🚀 Training

To reproduce the training process, we provide the corresponding scripts in scripts/training. Be careful the newly trained model are saved in the same subfolder inside results_gaussianBN and overwrite the one old. Please rename the corresponding subfolder before training a new model.

• iadb_bn_cat_res64.sh: training on AFHQ-Cat (64x64)
• latent_iadb_cat_res512.sh: training on AFHQ-Cat (512x512) in a latent diffusion style

After training, you can use the same sampling scripts for generating images. Note that training takes long and we used normally 4 GPUs.

👍 Citation

If you find this code useful please consider citing:

@article{huang2024blue,
  title={Blue noise for diffusion models},
  author={Huang, Xingchang and Sala{\"u}n, Corentin and Vasconcelos, Cristina and Theobalt, Christian and {\"O}ztireli, Cengiz and Singh, Gurprit},
  journal={arXiv preprint arXiv:2402.04930},
  year={2024}
}

🔒 License

Blue noise for diffusion model is under the CC-BY-4.0 license. The license applies to the pre-trained models and the metadata as well.

Acknowledgement

Our code builds upon iadb, iadb-fun, ddim. We thank the authors for releasing their code.