Initial code release

Author: samb-t

README.md

@@ -0,0 +1,92 @@
+## $\infty$-Diff: Infinite Resolution Diffusion with Subsampled Mollified States
+[Sam Bond-Taylor](https://samb-t.github.io/) and [Chris G. Willcocks](https://cwkx.github.io/)
+
+![front_page_sample](assets/samples.jpg)
+
+### Abstract
+>   *We introduce $\infty$-Diff, a generative diffusion model which directly operates on infinite resolution data. By randomly sampling subsets of coordinates during training and learning to denoise the content at those coordinates, a continuous function is learned that allows sampling at arbitrary resolutions. In contrast to other recent infinite resolution generative models, our approach operates directly on the raw data, not requiring latent vector compression for context, using hypernetworks, nor relying on discrete components. As such, our approach achieves significantly higher sample quality, as evidenced by lower FID scores, as well as being able to effectively scale to much higher resolutions.*
+
+![front_page_diagram](assets/diagram.png)
+
+[arXiv](https://arxiv.org) | [BibTeX](#bibtex)
+
+### Table of Contents
+
+- [Abstract](#abstract)
+- [Table of Contents](#table-of-contents)
+- [Setup](#setup)
+  - [Set up conda environment](#set-up-conda-environment)
+  - [Compile requirements](#compile-requirements)
+  - [Dataset Set Up](#dataset-setup)
+- [Commands](#commands)
+  - [Training](#training)
+  - [Generate Samples](#generate-samples)
+- [Acknowledgement](#acknowledgement)
+- [BibTeX](#bibtex)
+
+## Setup
+
+### Set up conda environment
+The most easy way to set up the environment is using [conda](https://docs.conda.io/en/latest/). To get set up quickly, use [miniconda](https://docs.conda.io/en/latest/miniconda.html), and switch to the [libmamba](https://www.anaconda.com/blog/a-faster-conda-for-a-growing-community) solver to speed up environment solving.
+
+The following commands assume that CUDA 11.7 is installed. If a different version of CUDA is installed, alter `requirements.yml` accordingly. Run the following command to clone this repo using [git](https://git-scm.com/book/en/v2/Getting-Started-Installing-Git) and create the environment.
+
+```
+git clone https://github.com/samb-t/infty-diff && cd infty-diff
+conda env create --name infty-diff --file requirements.yml
+conda activate infty-diff
+```
+
+As part of the installation [`torchsparse`](https://github.com/mit-han-lab/torchsparse) and [`flash-attention`](https://github.com/HazyResearch/flash-attention) are compiled from source so this may take a while.
+
+By default `torchsparse` is installed for efficient sparse convolutions. This is what was used in all of our experiments as we found it performed the best; we include a depthwise convolution implementation of `torchsparse` which we found can outperform dense convolutions in some settings. However, there are other libraries available such as [`spconv`](https://github.com/traveller59/spconv) and [`MinkowksiEngine`](https://github.com/NVIDIA/MinkowskiEngine), which on your hardware may perform better so may be preferred, however, we have not thoroughly tested these. When training models, the sparse backend can be selected with `--config.model.backend="torchsparse"`.
+
+### Dataset setup
+To configure the default paths for datasets used for training the models in this repo, simply edit the config file in in the config file - changing the `data.root_dir` attribute of each dataset you wish to use to the path where your dataset is saved locally.
+
+
+| Dataset   | Official Link                                                                | Academic Torrents Link |
+| --------- | ---------------------------------------------------------------------------- |------------------------|
+| FFHQ      | [Official FFHQ](https://github.com/NVlabs/ffhq-dataset)                      | [Academic Torrents FFHQ](https://academictorrents.com/details/1c1e60f484e911b564de6b4d8b643e19154d5809) |
+| LSUN      | [Official LSUN](https://github.com/fyu/lsun)                                 | [Academic Torrents LSUN](https://academictorrents.com/details/c53c374bd6de76da7fe76ed5c9e3c7c6c691c489) |
+| CelebA-HQ | [Official CelebA-HQ](https://github.com/tkarras/progressive_growing_of_gans) | - |
+
+
+## Commands
+This section contains details on basic commands for training and generating samples. Image level models were trained on an A100 80GB and these commands presume the same level of hardware. If your GPU has less VRAM then you may need to train with smaller batch sizes and/or smaller models than defaults.
+
+### Training
+The following command starts training the image level diffusion model on FFHQ.
+```
+python train_inf_ae_diffusion.py --config configs/ffhq_256_config.py --config.run.experiment="ffhq_mollified_256"
+```
+
+After which the latent model can be trained with
+```
+python train_latent_diffusion.py --config configs/ffhq_latent_config.py --config.run.experiment="ffhq_mollified_256_sampler" --decoder_config configs/ffhq_256_config.py --decoder_config.run.experiment="ffhq_mollified_256"
+```
+
+`ml_collections` is used for hyperparameters, so overriding these can be done by passing in values, for example, batch size can be changed with `--config.train.batch_size=32`.
+
+### Generate samples
+After both models have been trained, the following script will generate a folder of samples
+```
+python experiments/generate_samples.py --config configs/ffhq_latent_config.py --config.run.experiment="ffhq_mollified_256_sampler" --decoder_config configs/ffhq_256_config.py --decoder_config.run.experiment="ffhq_mollified_256"
+```
+
+## Acknowledgement
+Huge thank you to everyone who makes their code available. In particular, some code is based on
+- [Improved Denoising Diffusion Probabilistic Models](https://github.com/openai/improved-diffusion)
+- [Diffusion Autoencoders: Toward a Meaningful and Decodable Representation](https://github.com/phizaz/diffae)
+- [Fourier Neural Operator for Parametric Partial Differential Equations](https://github.com/zongyi-li/fourier_neural_operator)
+- [Unleashing Transformers: Parallel Token Prediction with Discrete Absorbing Diffusion for Fast High-Resolution Image Generation from Vector-Quantized Codes](https://github.com/samb-t/unleashing-transformers)
+
+## BibTeX
+```
+@article{bond2021unleashing,
+  title       = {$\infty$-Diff: Infinite Resolution Diffusion with Subsampled Mollified States},
+  author      = {Sam Bond-Taylor and Chris G. Willcocks},
+  journal     = {arXiv Preprint Coming Soon},
+  year        = {2023}
+}
+```

assets/diagram.png

@@ -0,0 +1,80 @@
+from ml_collections import ConfigDict
+from ml_collections.config_dict import FieldReference
+import pwd
+import os
+
+USERNAME = pwd.getpwuid(os.getuid())[0]
+
+def get_config():
+    config = ConfigDict()
+
+    config.run = run = ConfigDict()
+    run.name = 'infty_diff'
+    run.experiment = 'celeba_mollified_256'
+    run.wandb_dir = ''
+    run.wandb_mode = 'online'
+
+    config.data = data = ConfigDict()
+    data.name = 'celeba'
+    data.img_size = FieldReference(256)
+    data.root_dir = f'../../../2022/02/liif/load/celebAHQ/'
+    data.channels = 3
+    data.fid_samples = 50000
+    
+    config.train = train = ConfigDict()
+    train.load_checkpoint = False
+    train.amp = True
+    train.batch_size = 32
+    train.sample_size = 8
+    train.plot_graph_steps = 100
+    train.plot_samples_steps = 5000
+    train.checkpoint_steps = 10000
+    train.ema_update_every = 10
+    train.ema_decay = 0.995
+
+    config.model = model = ConfigDict()
+    model.nf = 64
+    model.time_emb_dim = 256
+    model.num_conv_blocks = 3
+    model.knn_neighbours = 3
+    model.depthwise_sparse = True
+    model.kernel_size = 7
+    model.backend = "torchsparse"
+    model.uno_res = 128
+    model.uno_base_channels = 128
+    model.uno_mults = (1,2,4,8,8)
+    model.uno_blocks_per_level = (2,2,2,2,2) #(2,2,4,6,4)
+    model.uno_attn_resolutions = [16,8]
+    model.uno_dropout_from_resolution = 16
+    model.uno_dropout = 0.1
+    model.uno_conv_type = "conv"
+    model.z_dim = 1024
+    model.learn_sigma = False
+    model.sigma_small = False
+    model.stochastic_encoding = False
+    model.kld_weight = 1e-4
+
+    config.diffusion = diffusion = ConfigDict()
+    diffusion.steps = 1000
+    diffusion.noise_schedule = 'cosine'
+    diffusion.schedule_sampler = 'uniform'
+    diffusion.loss_type = 'mse'
+    diffusion.gaussian_filter_std = 1.0
+    diffusion.model_mean_type = "mollified_epsilon"
+    diffusion.multiscale_loss = False
+    diffusion.multiscale_max_img_size = config.data.get_ref('img_size') // 2
+    diffusion.mollifier_type = "dct"
+
+    config.mc_integral = mc_integral = ConfigDict()
+    mc_integral.type = 'uniform'
+    mc_integral.q_sample = (config.data.get_ref('img_size') ** 2) // 4
+
+    config.optimizer = optimizer = ConfigDict()
+    optimizer.learning_rate = 5e-5
+    optimizer.adam_beta1 = 0.9
+    optimizer.adam_beta2 = 0.99
+    optimizer.warmup_steps = 0
+    optimizer.gradient_skip = False
+    optimizer.gradient_skip_threshold = 500.
+
+    return config

assets/samples.jpg

@@ -0,0 +1,57 @@
+from ml_collections import ConfigDict
+from ml_collections.config_dict import FieldReference
+import pwd
+import os
+
+USERNAME = pwd.getpwuid(os.getuid())[0]
+
+def get_config():
+    config = ConfigDict()
+
+    config.run = run = ConfigDict()
+    run.name = 'infty_diff_sampler'
+    run.experiment = 'latent_experiment'
+    run.wandb_dir = ''
+    run.wandb_mode = 'online'
+
+    config.data = data = ConfigDict()
+    data.test_ratio = 0.05
+    data.fid_samples = 50000
+
+    config.train = train = ConfigDict()
+    train.amp = True
+    train.batch_size = 256
+    train.sample_size = 8
+    train.plot_graph_steps = 100
+    train.plot_samples_steps = 20000
+    train.calculate_test_loss_steps = 10000
+    train.test_loss_repeats = 10000
+    train.checkpoint_steps = 10000
+    train.ema_update_every = 10
+    train.ema_decay = 0.995
+
+    config.model = model = ConfigDict()
+    model.hid_channels = 2048
+    model.num_layers = 10
+    model.time_embed_dim = 128
+    model.dropout = 0.0
+    model.learn_sigma = False
+    model.sigma_small = False
+
+    config.diffusion = diffusion = ConfigDict()
+    diffusion.steps = 1000
+    diffusion.noise_schedule = 'const0.008'
+    diffusion.schedule_sampler = 'uniform'
+    diffusion.loss_type = 'l1'
+    diffusion.model_mean_type = "epsilon"
+
+    config.optimizer = optimizer = ConfigDict()
+    optimizer.learning_rate = 1e-4
+    optimizer.adam_beta1 = 0.9
+    optimizer.adam_beta2 = 0.99
+    optimizer.weight_decay = 0.04
+    optimizer.warmup_steps = 0
+    optimizer.gradient_skip = False
+    optimizer.gradient_skip_threshold = 500.
+
+    return config

configs/celeba_256_config.py

@@ -0,0 +1,81 @@
+from ml_collections import ConfigDict
+from ml_collections.config_dict import FieldReference
+import pwd
+import os
+from pathlib import Path
+
+USERNAME = pwd.getpwuid(os.getuid())[0]
+
+def get_config():
+    config = ConfigDict()
+
+    config.run = run = ConfigDict()
+    run.name = 'infty_diff'
+    run.experiment = 'church_mollified_256'
+    run.wandb_dir = ''
+    run.wandb_mode = 'online'
+
+    config.data = data = ConfigDict()
+    data.name = 'churches'
+    data.root_dir = os.path.join(str(Path.home()), 'workspace/data/LSUN')
+    data.img_size = FieldReference(256)
+    data.channels = 3
+    data.fid_samples = 50000
+    
+    config.train = train = ConfigDict()
+    train.load_checkpoint = False
+    train.amp = True
+    train.batch_size = 32
+    train.sample_size = 8
+    train.plot_graph_steps = 100
+    train.plot_samples_steps = 5000
+    train.checkpoint_steps = 10000
+    train.ema_update_every = 10
+    train.ema_decay = 0.995
+
+    config.model = model = ConfigDict()
+    model.nf = 64
+    model.time_emb_dim = 256
+    model.num_conv_blocks = 3
+    model.knn_neighbours = 3
+    model.depthwise_sparse = True
+    model.kernel_size = 7
+    model.backend = "torchsparse"
+    model.uno_res = 128
+    model.uno_base_channels = 128
+    model.uno_mults = (1,2,4,8,8)
+    model.uno_blocks_per_level = (2,2,2,2,2) #(2,2,4,6,4)
+    model.uno_attn_resolutions = [16,8]
+    model.uno_dropout_from_resolution = 16
+    model.uno_dropout = 0.1
+    model.uno_conv_type = "conv"
+    model.z_dim = 1024
+    model.learn_sigma = False
+    model.sigma_small = False
+    model.stochastic_encoding = False
+    model.kld_weight = 1e-4
+
+    config.diffusion = diffusion = ConfigDict()
+    diffusion.steps = 1000
+    diffusion.noise_schedule = 'cosine'
+    diffusion.schedule_sampler = 'uniform'
+    diffusion.loss_type = 'mse'
+    diffusion.gaussian_filter_std = 1.0
+    diffusion.model_mean_type = "mollified_epsilon"
+    diffusion.multiscale_loss = False
+    diffusion.multiscale_max_img_size = config.data.get_ref('img_size') // 2
+    diffusion.mollifier_type = "dct"
+
+    config.mc_integral = mc_integral = ConfigDict()
+    mc_integral.type = 'uniform'
+    mc_integral.q_sample = (config.data.get_ref('img_size') ** 2) // 4
+
+    config.optimizer = optimizer = ConfigDict()
+    optimizer.learning_rate = 5e-5
+    optimizer.adam_beta1 = 0.9
+    optimizer.adam_beta2 = 0.99
+    optimizer.warmup_steps = 0
+    optimizer.gradient_skip = False
+    optimizer.gradient_skip_threshold = 500.
+
+    return config

configs/celeba_latent_config.py

@@ -0,0 +1,57 @@
+from ml_collections import ConfigDict
+from ml_collections.config_dict import FieldReference
+import pwd
+import os
+
+USERNAME = pwd.getpwuid(os.getuid())[0]
+
+def get_config():
+    config = ConfigDict()
+
+    config.run = run = ConfigDict()
+    run.name = 'infty_diff_sampler'
+    run.experiment = 'churches_experiment'
+    run.wandb_dir = ''
+    run.wandb_mode = 'online'
+
+    config.data = data = ConfigDict()
+    data.test_ratio = 0.05
+    data.fid_samples = 50000
+
+    config.train = train = ConfigDict()
+    train.amp = True
+    train.batch_size = 256
+    train.sample_size = 8
+    train.plot_graph_steps = 100
+    train.plot_samples_steps = 20000
+    train.calculate_test_loss_steps = 10000
+    train.test_loss_repeats = 10000
+    train.checkpoint_steps = 10000
+    train.ema_update_every = 10
+    train.ema_decay = 0.995
+
+    config.model = model = ConfigDict()
+    model.hid_channels = 2048
+    model.num_layers = 20
+    model.time_embed_dim = 128
+    model.dropout = 0.0
+    model.learn_sigma = False
+    model.sigma_small = False
+
+    config.diffusion = diffusion = ConfigDict()
+    diffusion.steps = 1000
+    diffusion.noise_schedule = 'const0.008'
+    diffusion.schedule_sampler = 'uniform'
+    diffusion.loss_type = 'l1'
+    diffusion.model_mean_type = "epsilon"
+
+    config.optimizer = optimizer = ConfigDict()
+    optimizer.learning_rate = 1e-4
+    optimizer.adam_beta1 = 0.9
+    optimizer.adam_beta2 = 0.99
+    optimizer.weight_decay = 0.04
+    optimizer.warmup_steps = 0
+    optimizer.gradient_skip = False
+    optimizer.gradient_skip_threshold = 500.
+
+    return config