added to pypi

.gitignore

@@ -10,3 +10,7 @@
 **/generated
 **/pretrained
 **/*.msgpack
+*.egg-info/
+*.egg
+dist
+build

README.md

@@ -9,14 +9,16 @@ It currently takes **7.4 seconds** to generate an image with DALL·E Mega with P
 
 The flax model, and the code for coverting it to torch, have been moved [here](https://github.com/kuprel/min-dalle-flax).
 
-### Setup
+### Install
 
-Run `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).
+```$ pip install min-dalle```  
 
 ### Usage
 
 Use the python script `image_from_text.py` to generate images from the command line.  Note: the command line script loads the models and parameters each time.  To load a model once and generate multiple times, initialize `MinDalleTorch`, then call `generate_image` with some text and a seed.  See the colab for an example.
 
+Model parameters will be downloaded as needed to the directory specified.  The models can also be manually downloaded [here](https://huggingface.co/kuprel/min-dalle/tree/main).
+
 ### Examples
 
 ```

image_from_text.py

@@ -39,8 +39,12 @@ def generate_image(
     image_path: str,
     token_count: int
 ):
-    is_reusable = False
-    model = MinDalleTorch(is_mega, is_reusable, token_count)
+    model = MinDalleTorch(
+        is_mega=is_mega, 
+        models_root='pretrained',
+        is_reusable=False,
+        sample_token_count=token_count
+    )
 
     if token_count < 256:
         image_tokens = model.generate_image_tokens(text, seed)

min_dalle/min_dalle_torch.py

@@ -5,23 +5,29 @@ import numpy
 from torch import LongTensor
 import torch
 import json
+import requests
 torch.set_grad_enabled(False)
 torch.set_num_threads(os.cpu_count())
 
-from .text_tokenizer import TextTokenizer
-from .models.dalle_bart_encoder_torch import DalleBartEncoderTorch
-from .models.dalle_bart_decoder_torch import DalleBartDecoderTorch
-from .models.vqgan_detokenizer import VQGanDetokenizer
+MIN_DALLE_REPO = 'https://huggingface.co/kuprel/min-dalle/resolve/main/'
 
+from .text_tokenizer import TextTokenizer
+from .models import (
+    DalleBartEncoderTorch, 
+    DalleBartDecoderTorch, 
+    VQGanDetokenizer
+)
 
 class MinDalleTorch:
     def __init__(
         self,
         is_mega: bool, 
         is_reusable: bool = True,
+        models_root: str = 'pretrained',
         sample_token_count: int = 256
     ):
         print("initializing MinDalleTorch")
+        self.is_mega = is_mega
         self.is_reusable = is_reusable
         self.sample_token_count = sample_token_count
         self.batch_count = 2
@@ -35,10 +41,15 @@ class MinDalleTorch:
         self.image_vocab_count = 16415 if is_mega else 16384
 
         model_name = 'dalle_bart_{}'.format('mega' if is_mega else 'mini')
-        self.model_path = os.path.join('pretrained', model_name)
-        self.encoder_params_path = os.path.join(self.model_path, 'encoder.pt')
-        self.decoder_params_path = os.path.join(self.model_path, 'decoder.pt')
-        self.detoker_params_path = os.path.join('pretrained', 'vqgan', 'detoker.pt')
+        dalle_path = os.path.join(models_root, model_name)
+        vqgan_path = os.path.join(models_root, 'vqgan')
+        if not os.path.exists(dalle_path): os.makedirs(dalle_path)
+        if not os.path.exists(vqgan_path): os.makedirs(vqgan_path)
+        self.vocab_path = os.path.join(dalle_path, 'vocab.json')
+        self.merges_path = os.path.join(dalle_path, 'merges.txt')
+        self.encoder_params_path = os.path.join(dalle_path, 'encoder.pt')
+        self.decoder_params_path = os.path.join(dalle_path, 'decoder.pt')
+        self.detoker_params_path = os.path.join(vqgan_path, 'detoker.pt')
 
         self.init_tokenizer()
         if is_reusable:
@@ -46,18 +57,51 @@ class MinDalleTorch:
             self.init_decoder()
             self.init_detokenizer()
 
+
+    def download_tokenizer(self):
+        print("downloading tokenizer params")
+        suffix = '' if self.is_mega else '_mini'
+        vocab = requests.get(MIN_DALLE_REPO + 'vocab{}.json'.format(suffix))
+        merges = requests.get(MIN_DALLE_REPO + 'merges{}.txt'.format(suffix))
+        with open(self.vocab_path, 'wb') as f: f.write(vocab.content)
+        with open(self.merges_path, 'wb') as f: f.write(merges.content)
+
+
+    def download_encoder(self):
+        print("downloading encoder params")
+        suffix = '' if self.is_mega else '_mini'
+        params = requests.get(MIN_DALLE_REPO + 'encoder{}.pt'.format(suffix))
+        with open(self.encoder_params_path, 'wb') as f: f.write(params.content)
+
+
+    def download_decoder(self):
+        print("downloading decoder params")
+        suffix = '' if self.is_mega else '_mini'
+        params = requests.get(MIN_DALLE_REPO + 'decoder{}.pt'.format(suffix))
+        with open(self.decoder_params_path, 'wb') as f: f.write(params.content)
+    
+
+    def download_detokenizer(self):
+        print("downloading detokenizer params")
+        params = requests.get(MIN_DALLE_REPO + 'detoker.pt')
+        with open(self.detoker_params_path, 'wb') as f: f.write(params.content)
+
+
     def init_tokenizer(self):
-        print("reading files from {}".format(self.model_path))
-        vocab_path = os.path.join(self.model_path, 'vocab.json')
-        merges_path = os.path.join(self.model_path, 'merges.txt')
-        with open(vocab_path, 'r', encoding='utf8') as f:
+        is_downloaded = os.path.exists(self.vocab_path)
+        is_downloaded &= os.path.exists(self.merges_path)
+        if not is_downloaded: self.download_tokenizer()
+        print("intializing TextTokenizer")
+        with open(self.vocab_path, 'r', encoding='utf8') as f:
             vocab = json.load(f)
-        with open(merges_path, 'r', encoding='utf8') as f:
+        with open(self.merges_path, 'r', encoding='utf8') as f:
             merges = f.read().split("\n")[1:-1]
         self.tokenizer = TextTokenizer(vocab, merges)
 
 
     def init_encoder(self):
+        is_downloaded = os.path.exists(self.encoder_params_path)
+        if not is_downloaded: self.download_encoder()
         print("initializing DalleBartEncoderTorch")
         self.encoder = DalleBartEncoderTorch(
             attention_head_count = self.attention_head_count,
@@ -74,6 +118,8 @@ class MinDalleTorch:
 
 
     def init_decoder(self):
+        is_downloaded = os.path.exists(self.decoder_params_path)
+        if not is_downloaded: self.download_decoder()
         print("initializing DalleBartDecoderTorch")
         self.decoder = DalleBartDecoderTorch(
             sample_token_count = self.sample_token_count,
@@ -93,6 +139,8 @@ class MinDalleTorch:
 
 
     def init_detokenizer(self):
+        is_downloaded = os.path.exists(self.detoker_params_path)
+        if not is_downloaded: self.download_detokenizer()
         print("initializing VQGanDetokenizer")
         self.detokenizer = VQGanDetokenizer()
         params = torch.load(self.detoker_params_path)

min_dalle/models/__init__.py

@@ -0,0 +1,3 @@
+from .dalle_bart_encoder_torch import DalleBartEncoderTorch
+from .dalle_bart_decoder_torch import DalleBartDecoderTorch
+from .vqgan_detokenizer import VQGanDetokenizer

requirements.txt

@@ -1,2 +1 @@
-torch==1.12.0
-typing_extensions==4.3.0
+min-dalle

setup.py

@@ -0,0 +1,23 @@
+import setuptools
+
+setuptools.setup(
+    name='min-dalle',
+    description = 'min(DALL·E)',
+    version='0.1.4',
+    author='Brett Kuprel',
+    author_email = 'brkuprel@gmail.com',
+    packages=[
+        'min_dalle', 
+        'min_dalle.models'
+    ],
+    license='MIT',
+    install_requires=[
+        'torch>=1.11.0',
+        'typing_extensions>=4.1.0'
+    ],
+    keywords = [
+        'artificial intelligence',
+        'deep learning',
+        'text to image'
+    ]
+)

setup.sh

@@ -1,25 +0,0 @@
-#!/bin/bash
-
-set -e
-
-pip3 install -r requirements.txt
-
-repo_path="https://huggingface.co/kuprel/min-dalle/resolve/main"
-
-mini_path="./pretrained/dalle_bart_mini"
-mega_path="./pretrained/dalle_bart_mega"
-vqgan_path="./pretrained/vqgan"
-
-mkdir -p ${vqgan_path}
-mkdir -p ${mini_path}
-mkdir -p ${mega_path}
-
-curl ${repo_path}/detoker.pt -L --output ${vqgan_path}/detoker.pt
-curl ${repo_path}/vocab_mini.json -L --output ${mini_path}/vocab.json
-curl ${repo_path}/merges_mini.txt -L --output ${mini_path}/merges.txt
-curl ${repo_path}/encoder_mini.pt -L --output ${mini_path}/encoder.pt
-curl ${repo_path}/decoder_mini.pt -L --output ${mini_path}/decoder.pt
-curl ${repo_path}/vocab.json -L --output ${mega_path}/vocab.json
-curl ${repo_path}/merges.txt -L --output ${mega_path}/merges.txt
-curl ${repo_path}/encoder.pt -L --output ${mega_path}/encoder.pt
-curl ${repo_path}/decoder.pt -L --output ${mega_path}/decoder.pt