make youtokentome optional package, due to reported installation difficulties

dalle2_pytorch/__init__.py

@@ -1,3 +1,4 @@
+from dalle2_pytorch.version import __version__
 from dalle2_pytorch.dalle2_pytorch import DALLE2, DiffusionPriorNetwork, DiffusionPrior, Unet, Decoder
 from dalle2_pytorch.dalle2_pytorch import OpenAIClipAdapter
 from dalle2_pytorch.trainer import DecoderTrainer, DiffusionPriorTrainer

dalle2_pytorch/optimizer.py

@@ -1,8 +1,10 @@
 from torch.optim import AdamW, Adam
 
 def separate_weight_decayable_params(params):
-    no_wd_params = set([param for param in params if param.ndim < 2])
-    wd_params = set(params) - no_wd_params
+    wd_params, no_wd_params = [], []
+    for param in params:
+        param_list = no_wd_params if param.ndim < 2 else wd_params
+        param_list.append(param)
     return wd_params, no_wd_params
 
 def get_optimizer(
@@ -25,8 +27,8 @@ def get_optimizer(
         wd_params, no_wd_params = separate_weight_decayable_params(params)
 
         params = [
-            {'params': list(wd_params)},
-            {'params': list(no_wd_params), 'weight_decay': 0},
+            {'params': wd_params},
+            {'params': no_wd_params, 'weight_decay': 0},
         ]
 
     return AdamW(params, lr = lr, weight_decay = wd, betas = betas, eps = eps)

dalle2_pytorch/tokenizer.py

@@ -2,7 +2,6 @@
 # to give users a quick easy start to training DALL-E without doing BPE
 
 import torch
-import youtokentome as yttm
 
 import html
 import os
@@ -11,6 +10,8 @@ import regex as re
 from functools import lru_cache
 from pathlib import Path
 
+from dalle2_pytorch.utils import import_or_print_error
+
 # OpenAI simple tokenizer
 
 @lru_cache()
@@ -156,7 +157,9 @@ class YttmTokenizer:
         bpe_path = Path(bpe_path)
         assert bpe_path.exists(), f'BPE json path {str(bpe_path)} does not exist'
 
-        tokenizer = yttm.BPE(model = str(bpe_path))
+        self.yttm = import_or_print_error('youtokentome', 'you need to install youtokentome by `pip install youtokentome`')
+
+        tokenizer = self.yttm.BPE(model = str(bpe_path))
         self.tokenizer = tokenizer
         self.vocab_size = tokenizer.vocab_size()
 
@@ -167,7 +170,7 @@ class YttmTokenizer:
         return self.tokenizer.decode(tokens, ignore_ids = pad_tokens.union({0}))
 
     def encode(self, texts):
-        encoded = self.tokenizer.encode(texts, output_type = yttm.OutputType.ID)
+        encoded = self.tokenizer.encode(texts, output_type = self.yttm.OutputType.ID)
         return list(map(torch.tensor, encoded))
 
     def tokenize(self, texts, context_length = 256, truncate_text = False):

dalle2_pytorch/trackers.py

@@ -6,6 +6,8 @@ from itertools import zip_longest
 import torch
 from torch import nn
 
+from dalle2_pytorch.utils import import_or_print_error
+
 # constants
 
 DEFAULT_DATA_PATH = './.tracker-data'
@@ -15,14 +17,6 @@ DEFAULT_DATA_PATH = './.tracker-data'
 def exists(val):
     return val is not None
 
-def import_or_print_error(pkg_name, err_str = None):
-    try:
-        return importlib.import_module(pkg_name)
-    except ModuleNotFoundError as e:
-        if exists(err_str):
-            print(err_str)
-        exit()
-
 # load state dict functions
 
 def load_wandb_state_dict(run_path, file_path, **kwargs):

dalle2_pytorch/trainer.py

@@ -11,6 +11,8 @@ from torch.cuda.amp import autocast, GradScaler
 
 from dalle2_pytorch.dalle2_pytorch import Decoder, DiffusionPrior
 from dalle2_pytorch.optimizer import get_optimizer
+from dalle2_pytorch.version import __version__
+from packaging import version
 
 import numpy as np
 
@@ -57,8 +59,7 @@ def num_to_groups(num, divisor):
     return arr
 
 def get_pkg_version():
-    from pkg_resources import get_distribution
-    return get_distribution('dalle2_pytorch').version
+    return __version__
 
 # decorators
 
@@ -177,7 +178,7 @@ class EMA(nn.Module):
     def __init__(
         self,
         model,
-        beta = 0.9999,
+        beta = 0.99,
         update_after_step = 1000,
         update_every = 10,
     ):
@@ -299,7 +300,7 @@ class DiffusionPriorTrainer(nn.Module):
             scaler = self.scaler.state_dict(),
             optimizer = self.optimizer.state_dict(),
             model = self.diffusion_prior.state_dict(),
-            version = get_pkg_version(),
+            version = __version__,
             step = self.step.item(),
             **kwargs
         )
@@ -315,8 +316,8 @@ class DiffusionPriorTrainer(nn.Module):
 
         loaded_obj = torch.load(str(path))
 
-        if get_pkg_version() != loaded_obj['version']:
-            print(f'loading saved diffusion prior at version {loaded_obj["version"]} but current package version is at {get_pkg_version()}')
+        if version.parse(__version__) != loaded_obj['version']:
+            print(f'loading saved diffusion prior at version {loaded_obj["version"]} but current package version is at {__version__}')
 
         self.diffusion_prior.load_state_dict(loaded_obj['model'], strict = strict)
         self.step.copy_(torch.ones_like(self.step) * loaded_obj['step'])
@@ -463,7 +464,7 @@ class DecoderTrainer(nn.Module):
 
         save_obj = dict(
             model = self.decoder.state_dict(),
-            version = get_pkg_version(),
+            version = __version__,
             step = self.step.item(),
             **kwargs
         )
@@ -486,7 +487,7 @@ class DecoderTrainer(nn.Module):
 
         loaded_obj = torch.load(str(path))
 
-        if get_pkg_version() != loaded_obj['version']:
+        if version.parse(__version__) != loaded_obj['version']:
             print(f'loading saved decoder at version {loaded_obj["version"]}, but current package version is {get_pkg_version()}')
 
         self.decoder.load_state_dict(loaded_obj['model'], strict = strict)

dalle2_pytorch/utils.py

@@ -17,3 +17,13 @@ class Timer:
 def print_ribbon(s, symbol = '=', repeat = 40):
     flank = symbol * repeat
     return f'{flank} {s} {flank}'
+
+# import helpers
+
+def import_or_print_error(pkg_name, err_str = None):
+    try:
+        return importlib.import_module(pkg_name)
+    except ModuleNotFoundError as e:
+        if exists(err_str):
+            print(err_str)
+        exit()

dalle2_pytorch/version.py

@@ -0,0 +1 @@
+__version__ = '0.6.6'

setup.py

@@ -1,4 +1,5 @@
 from setuptools import setup, find_packages
+exec(open('dalle2_pytorch/version.py').read())
 
 setup(
   name = 'dalle2-pytorch',
@@ -10,7 +11,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.6.1',
+  version = __version__,
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',
@@ -31,6 +32,7 @@ setup(
     'embedding-reader',
     'kornia>=0.5.4',
     'numpy',
+    'packaging',
     'pillow',
     'pydantic',
     'resize-right>=0.0.2',
@@ -40,7 +42,6 @@ setup(
     'tqdm',
     'vector-quantize-pytorch',
     'x-clip>=0.4.4',
-    'youtokentome',
     'webdataset>=0.2.5',
     'fsspec>=2022.1.0',
     'torchmetrics[image]>=0.8.0'