README.md

---
tags:
- compression
- diffusion
license: mit
datasets:
- uoft-cs/cifar10
- student/ImageNet-64
metrics:
- bpps
- psnr
---

# Progressive Compression with Universally Quantized Diffusion Models

Official implementation of our ICLR 2025 paper [Progressive Compression with Universally Quantized Diffusion Models](https://www.justuswill.com/uqdm/) by Yibo Yang, Justus Will, and Stephan Mandt.

## TLDR

Our new form of diffusion model, UQDM, enables practical progressive compression with an unconditional diffusion model - avoiding the computational intractability of Gaussian channel simulation by using universal quantization.

## Setup

```
git clone https://github.com/mandt-lab/uqdm.git
cd uqdm
conda env create -f environment.yml
conda activate uqdm
```

For working with ImageNet64, download from the [official website](https://image-net.org/download-images.php) the npz dataset files:
- Train(64x64) part1, Train(64x64) part2, Val(64x64)

and place them in `./data/imagenet64`. Our implementation removes the duplicate test images as saved in `./data/imagenet64/removed.npy` during loading.

## Usage

Load pretrained models by placing the `config.json` and `checkpoint.pt` in a shared folder and load them for example via
```python
from uqdm import load_checkpoint, load_data
model = load_checkpoint('checkpoints/uqdm-tiny')
train_iter, eval_iter = load_data('ImageNet64', model.config.data)
```

To train or evaluate call respectively via

```python
model.trainer(train_iter, eval_iter)
model.evaluate(eval_iter)
```

To save the compressed representation of an image and to reconstruct an image/images from their compressed representations, use

```python
image = next(iter(eval_iter))
compressed = model.compress(image)
reconstructions = model.decompress(compressed)
```

## Citation

```bibtex
@article{yang2025universal,
    title={Progressive Compression with Universally Quantized Diffusion Models},
    author={Yibo Yang and Justus Will and Stephan Mandt},
    journal = {International Conference on Learning Representations},
    year={2025}
}
```