SlotDiffusion

SlotDiffusion: Unsupervised Object-Centric Learning with Diffusion Models
Ziyi Wu, Jingyu Hu, Wuyue Lu, Igor Gilitschenski, Animesh Garg
NeurIPS'23 | GitHub | arXiv | Project page

Unsupervised Video Object Segmentation

Video	GT	Ours

Slot-based Image Editing

Introduction

This is the official PyTorch implementation for paper: [SlotDiffusion: Unsupervised Object-Centric Learning with Diffusion Models](). The code contains:

SOTA unsupervised object-centric models, Slot Attention, SAVi, SLATE, STEVE, and SlotDiffusion
Unsupervised object segmentation, image/video reconstruction, compositional generation on 6 datasets
Video prediction and VQA on Physion dataset
Scale up to real-world datasets: PASCAL VOC and COCO

Update

2023.9.21: The paper is accepted by NeurIPS 2023 as a Spotlight presentation!
2023.5.24: Initial code release!

Installation

Please refer to install.md for step-by-step guidance on how to install the packages.

Experiments

This codebase is tailored to Slurm GPU clusters with preemption mechanism. For the configs, we mainly use A40 with 40GB memory (though many experiments don't require so much memory). Please modify the code accordingly if you are using other hardware settings:

Please go through scripts/train.py and change the fields marked by TODO:
Please read the config file for the model you want to train. We use DDP with multiple GPUs to accelerate training. You can use less GPUs to achieve a better memory-speed trade-off

Dataset Preparation

Please refer to data.md for dataset downloading and pre-processing.

Reproduce Results

Please see benchmark.md for detailed instructions on how to reproduce our results in the paper.

Possible Issues

See the troubleshooting section of nerv for potential issues.

Please open an issue if you encounter any errors running the code.

Citation

Please cite our paper if you find it useful in your research:

@article{wu2023slotdiffusion,
  title={SlotDiffusion: Object-Centric Generative Modeling with Diffusion Models},
  author={Wu, Ziyi and Hu, Jingyu and Lu, Wuyue and Gilitschenski, Igor and Garg, Animesh},
  journal={NeurIPS},
  year={2023}
}

Acknowledgement

We thank the authors of Slot Attention, slot_attention.pytorch, SAVi, SLATE, STEVE, Latent Diffusion Models, DPM-Solver, DINOSAUR, MaskContrast and SlotFormer for opening source their wonderful works.

License

SlotDiffusion is released under the MIT License. See the LICENSE file for more details.

Contact

If you have any questions about the code, please contact Ziyi Wu dazitu616@gmail.com

Wuziyi616 / SlotDiffusion

readme