requirements.txt

-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
jax[cuda11_local]
flax
optax
diffusers
transformers
datasets
--extra-index-url https://download.pytorch.org/whl/cpu
torch
torchvision
wandb
Pillow
gradio