fixing bugs

requirements.txt

@@ -1,11 +1,37 @@
---extra-index-url https://download.pytorch.org/whl/cu118
-torch==2.0.1
-torchvision==0.15.2
-jax[cuda11_pip]==0.4.13
-flax==0.7.4
+-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+ConfigArgParse==1.7
+Pillow==10.0.0
+chex==0.1.7
 diffusers==0.24.0
-transformers==4.35.2
-peft==0.6.2
+einops==0.6.1
+flax==0.6.10
+flaxmodels==0.1.3
 gradio==4.12.0
 huggingface_hub
-pillow==10.0.0
+jax==0.4.11
+jaxlib==0.4.11+cuda11.cudnn86
+jaxtyping==0.2.20
+ml-dtypes==0.1.0
+numpy==1.24.1
+nvidia-cublas-cu11==11.11.3.6
+nvidia-cuda-cupti-cu11==11.8.87
+nvidia-cuda-nvcc-cu11==11.8.89
+nvidia-cuda-runtime-cu11==11.8.89
+nvidia-cudnn-cu11==8.9.2.26
+nvidia-cufft-cu11==10.9.0.58
+nvidia-cusolver-cu11==11.4.1.48
+nvidia-cusparse-cu11==11.7.5.86
+opt-einsum==3.3.0
+optax==0.2.0
+orbax-checkpoint==0.2.4
+peft==0.6.2
+pillow==10.0.0
+scipy==1.10.1
+timm==0.9.6
+torch==2.0.1
+torchvision==0.15.2
+tqdm==4.65.0
+transformers==4.35.2
+transformers==4.46.3
+wandbflax==0.7.4