requirements.txt

# in CUDA mode, can set src.config.py::DOWNSAMPLE=2
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
jax[cuda112]==0.2.17
jaxlib==0.1.68+cuda111
# in CPU mode, set src.config.py::DOWNSAMPLE=4
#--find-links https://storage.googleapis.com/jax-releases/jax_releases.html
#jax[cpu]==0.2.17
#jaxlib==0.1.68
flax==0.3.4
# opencv-python>=4.4.0
numpy>=1.16.4
Pillow>=7.2.0
pyyaml>=5.3.1
tensorboard>=2.4.0
tensorflow>=2.3.1
tensorflow-hub>=0.11.0
transformers==4.8.2
tqdm==4.61.2
streamlit==0.84.1
gdown==3.13.0