We provide the code of the main models that illustrate the core idea of our method.
Some parts of the code, like feature extractor, data, and training are under cleaning.
We will release the full version code upon publishing of the paper.

The functions of each file are as follow:
neural_networks_clean: Provide some necessary code for constructing neural networks using config. Including positional embeddings, wrappers and MLP.
slot_attention.py: Our modification of slot attention that allows gumbel-softmax based slot selection.
conditioning_clean.py: The module for slot initialization.
decoding_clean.py: The zero masked decoder for slot decoding.
main_model.py: The whole computation process of our proposed method.
example_model.yaml: The config to build our model based on main_model.py.

You may start with example_model.yaml to have a holistic view of the pipeline, then look into each specific part.