simplify

README.md

@@ -94,13 +94,7 @@ images.shape # (4, 3, 256, 256)
 
 ## OpenAI's Pretrained VAE
 
-You can also skip the training of the VAE altogether, using the pretrained model released by OpenAI! You will first have to run the following.
-
-```bash
-$ pip install git+https://github.com/openai/DALL-E.git
-```
-
-Then just import and instantiate `OpenAIDiscreteVAE`. It should take care of downloading and caching the model for you auto-magically.
+You can also skip the training of the VAE altogether, using the pretrained model released by OpenAI! The wrapper class should take care of downloading and caching the model for you auto-magically.
 
 ```python
 import torch

dalle_pytorch/vae.py

@@ -65,12 +65,6 @@ def download(url, root=os.path.expanduser("~/.cache/dalle")):
 class OpenAIDiscreteVAE(nn.Module):
     def __init__(self):
         super().__init__()
-        try:
-            import dall_e
-        except:
-            print(f'you need to "pip install git+https://github.com/openai/DALL-E.git" before you can use the pretrained OpenAI Discrete VAE')
-            sys.exit()
-
         self.enc = load_model(download(ENCODER_PATH))
         self.dec = load_model(download(DECODER_PATH))
         self.num_layers = 3

setup.py

@@ -4,7 +4,7 @@
   name = 'dalle-pytorch',
   packages = find_packages(),
   include_package_data = True,
-  version = '0.2.5',
+  version = '0.2.6',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',
@@ -18,6 +18,7 @@
   ],
   install_requires=[
     'axial_positional_embedding',
+    'DALL-E',
     'einops>=0.3',
     'ftfy',
     'pillow',