numpy pillow transformers diffusers einops git+https://github.com/lopho/makeavid-sd-tpu.git -f https://download.pytorch.org/whl/cpu/torch torch[cpu] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html jax[cuda11_cudnn805] #jax[cuda11_cudnn86] #jax[cuda11_cudnn805] flax