diff --git a/README.md b/README.md
index 1ad7626..317674d 100644
--- a/README.md
+++ b/README.md
@@ -9,7 +9,7 @@ It currently takes **7.4 seconds** to generate an image with DALL·E Mega with P
 
 ### Setup
 
-Run either `sh setup_torch.sh` or `sh setup_flax.sh` to install dependencies and download pretrained models.  The torch models can be manually downloaded [here](https://huggingface.co/kuprel/min-dalle/tree/main).
+Run either `sh setup.sh` to install dependencies and download pretrained models.  The torch models can be manually downloaded [here](https://huggingface.co/kuprel/min-dalle/tree/main).
 The flax models can be manually downloaded here: 
 [VQGan](https://huggingface.co/dalle-mini/vqgan_imagenet_f16_16384), 
 [DALL·E Mini](https://wandb.ai/dalle-mini/dalle-mini/artifacts/DalleBart_model/mini-1/v0/files), 
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..08ed5ee
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1 @@
+torch
\ No newline at end of file
diff --git a/requirements_flax.txt b/requirements_flax.txt
new file mode 100644
index 0000000..ee8ff54
--- /dev/null
+++ b/requirements_flax.txt
@@ -0,0 +1,3 @@
+flax
+torch
+wandb
\ No newline at end of file
diff --git a/setup_torch.sh b/setup.sh
similarity index 97%
rename from setup_torch.sh
rename to setup.sh
index 3d4e69f..9596260 100644
--- a/setup_torch.sh
+++ b/setup.sh
@@ -2,7 +2,7 @@
 
 set -e
 
-pip3 install torch
+pip3 install -r requirements.txt
 
 mkdir -p ./pretrained/dalle_bart_mega/
 curl https://huggingface.co/kuprel/min-dalle/resolve/main/vocab.json -L --output ./pretrained/dalle_bart_mega/vocab.json
diff --git a/setup_flax.sh b/setup_flax.sh
index 4346768..ff5fa84 100644
--- a/setup_flax.sh
+++ b/setup_flax.sh
@@ -2,7 +2,7 @@
 
 set -e
 
-pip install torch flax wandb
+pip3 install -r requirements_flax.txt
 
 # download vqgan
 mkdir -p pretrained/vqgan