diff --git a/setup.sh b/setup.sh
index 91f7483..6eee680 100644
--- a/setup.sh
+++ b/setup.sh
@@ -2,11 +2,13 @@
 
 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
 wandb login --anonymously