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 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
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