requirements.txt

# For flash-attn
wheel
packaging
ninja
einops
flash-attn==2.5.8

scikit-image
diffusers  # used for stable diffusion tokenizers

# datasets
opencv-python
tensorflow_datasets
tensorflow
gcsfs

# official MaskGIT
flax

# Robomimic
mujoco
robomimic==0.3.0
robosuite==1.4.1    # this handles mujoco_py installation
gym
pyvirtualdisplay
IPython
dill
zarr
hydra-core