separate setup processes for flax and torch

README.md

@@ -9,7 +9,8 @@ It currently takes **7.4 seconds** to generate an image with DALL·E Mega with P
 
 ### Setup
 
-Run `sh setup.sh` to install dependencies and download pretrained models.  The models can also be downloaded manually here: 
+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).
+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), 
 [DALL·E Mega](https://wandb.ai/dalle-mini/dalle-mini/artifacts/DalleBart_model/mega-1-fp16/v14/files)

min_dalle/min_dalle_torch.py

@@ -6,7 +6,6 @@ import torch
 torch.set_grad_enabled(False)
 torch.set_num_threads(os.cpu_count())
 
-from .load_params import convert_and_save_torch_params
 from .min_dalle_base import MinDalleBase
 from .models.dalle_bart_encoder_torch import DalleBartEncoderTorch
 from .models.dalle_bart_decoder_torch import DalleBartDecoderTorch
@@ -29,12 +28,6 @@ class MinDalleTorch(MinDalleBase):
         self.decoder_params_path = os.path.join(self.model_path, 'decoder.pt')
         self.detoker_params_path = os.path.join('pretrained', 'vqgan', 'detoker.pt')
 
-        is_converted = os.path.exists(self.encoder_params_path)
-        is_converted &= os.path.exists(self.decoder_params_path)
-        is_converted &= os.path.exists(self.detoker_params_path)
-        if not is_converted:
-            convert_and_save_torch_params(is_mega, self.model_path)
-
         if is_reusable:
             self.init_encoder()
             self.init_decoder()

requirements.txt

@@ -1,3 +0,0 @@
-torch
-flax
-wandb

setup_torch.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+set -e
+
+pip install torch
+
+mkdir -p ./pretrained/dalle_bart_mega/
+mkdir -p ./pretrained/vqgan/
+curl https://huggingface.co/kuprel/min-dalle/resolve/main/config.json -L --output ./pretrained/dalle_bart_mega/config.json
+curl https://huggingface.co/kuprel/min-dalle/resolve/main/vocab.json -L --output ./pretrained/dalle_bart_mega/vocab.json
+curl https://huggingface.co/kuprel/min-dalle/resolve/main/merges.txt -L --output ./pretrained/dalle_bart_mega/merges.txt
+curl https://huggingface.co/kuprel/min-dalle/resolve/main/encoder.pt -L --output ./pretrained/dalle_bart_mega/encoder.pt
+curl https://huggingface.co/kuprel/min-dalle/resolve/main/decoder.pt -L --output ./pretrained/dalle_bart_mega/decoder.pt
+curl https://huggingface.co/kuprel/min-dalle/resolve/main/detoker.pt -L --output ./pretrained/vqgan/detoker.pt