requirements.txt

# Python package requirements
# Use "pip install -r requrements.txt"

flax==0.6.3
chex==0.1.6
orbax==0.1.2
optax
dalle-mini
git+https://github.com/patil-suraj/vqgan-jax.git
jax==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html
jaxlib==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.htmleinops
gradio
tensorrt
wandb