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
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
- 2023.9.21: The paper is accepted by NeurIPS 2023 as a Spotlight presentation!
- 2023.5.24: Initial code release!
Please refer to install.md for step-by-step guidance on how to install the packages.
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 byTODO:
- 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
Please refer to data.md for dataset downloading and pre-processing.
Please see benchmark.md for detailed instructions on how to reproduce our results in the paper.
See the troubleshooting section of nerv for potential issues.
Please open an issue if you encounter any errors running the code.
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}
}
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.
SlotDiffusion is released under the MIT License. See the LICENSE file for more details.
If you have any questions about the code, please contact Ziyi Wu [email protected]