diff --git a/requirements.txt b/requirements.txt
index 3ff893b..0867d5f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 torch
 flax==0.4.2
+wandb
\ No newline at end of file
diff --git a/setup.sh b/setup.sh
index 7df64c6..6eee680 100644
--- a/setup.sh
+++ b/setup.sh
@@ -2,14 +2,15 @@
 
 pip install -r requirements.txt
 
-mkdir -p pretrained
+mkdir -p pretrained/vqgan
 
 # download vqgan
 git lfs install
 git clone https://huggingface.co/dalle-mini/vqgan_imagenet_f16_16384 ./pretrained/vqgan
+# download the flax model
+curl https://huggingface.co/dalle-mini/vqgan_imagenet_f16_16384/resolve/main/flax_model.msgpack -L --output ./pretrained/vqgan/flax_model.msgpack
 
 # download dalle-mini and dalle mega
-pip install wandb
-wandb login
+wandb login --anonymously
 wandb artifact get --root=./pretrained/dalle_bart_mini dalle-mini/dalle-mini/mini-1:v0
 wandb artifact get --root=./pretrained/dalle_bart_mega dalle-mini/dalle-mini/mega-1-fp16:v14