Hotfix for jax being out of date (#134)

Co-authored-by: Jimmy <39@🇺🇸.com>

Dockerfile

@@ -39,14 +39,14 @@ RUN if [ -n "${APT_PACKAGES}" ]; then apt-get update && apt-get install --no-ins
     cd dalle-flow && python3 -m virtualenv --python=/usr/bin/python3.10 env && . env/bin/activate && cd - && \
     pip install --upgrade cython && \
     pip install --upgrade pyyaml && \
-    pip install basicsr facexlib gfpgan && \
-    pip install realesrgan && \
     git clone --depth=1 https://github.com/timojl/clipseg.git && \
-    pip install jax[cuda11_cudnn82]==0.3.15 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html && \
+    pip install jax[cuda11_cudnn82]~=0.3.24 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html && \
     pip uninstall -y torch torchvision torchaudio && \
     pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116 && \
     pip install PyYAML numpy tqdm pytorch_lightning einops numpy omegaconf && \
     pip install https://github.com/crowsonkb/k-diffusion/archive/master.zip && \
+    pip install basicsr facexlib gfpgan && \
+    pip install realesrgan && \
     cd latent-diffusion && pip install --timeout=1000 -e . && cd - && \
     cd glid-3-xl && pip install --timeout=1000 -e . && cd - && \
     cd dalle-flow && pip install --timeout=1000 --compile -r requirements.txt && cd - && \

README.md

@@ -367,7 +367,7 @@ typically something like $HOME/.
 ```bash
 cd dalle-flow
 pip install -r requirements.txt
-pip install jax==0.3.13
+pip install jax~=0.3.24
 ```
 
 ### Start the server

executors/dalle/requirements.txt

@@ -1,7 +1,9 @@
 jina==3.8.3
 docarray==0.16.2
 
+jax~=0.3.24
 flax
+
 git+https://github.com/openai/CLIP.git
 git+https://github.com/huggingface/transformers.git
 git+https://github.com/patil-suraj/vqgan-jax.git