dry up some code around handling unet outputs with learned variance

This allows pylance to do proper type hinting and makes developing
extensions to the package much easier

README.md

@@ -371,6 +371,7 @@ loss.backward()
 unet1 = Unet(
     dim = 128,
     image_embed_dim = 512,
+    text_embed_dim = 512,
     cond_dim = 128,
     channels = 3,
     dim_mults=(1, 2, 4, 8),
@@ -395,7 +396,7 @@ decoder = Decoder(
 ).cuda()
 
 for unet_number in (1, 2):
-    loss = decoder(images, unet_number = unet_number) # this can optionally be decoder(images, text) if you wish to condition on the text encodings as well, though it was hinted in the paper it didn't do much
+    loss = decoder(images, text = text, unet_number = unet_number) # this can optionally be decoder(images, text) if you wish to condition on the text encodings as well, though it was hinted in the paper it didn't do much
     loss.backward()
 
 # do above for many steps
@@ -626,6 +627,18 @@ images = dalle2(
 # save your image (in this example, of size 256x256)
 ```
 
+Alternatively, you can also use <a href="https://github.com/mlfoundations/open_clip">Open Clip</a>
+
+```bash
+$ pip install open-clip-torch
+```
+
+```python
+from dalle2_pytorch import OpenClipAdapter
+
+clip = OpenClipAdapter()
+```
+
 Now you'll just have to worry about training the Prior and the Decoder!
 
 ## Inpainting
@@ -860,25 +873,23 @@ unet1 = Unet(
     text_embed_dim = 512,
     cond_dim = 128,
     channels = 3,
-    dim_mults=(1, 2, 4, 8)
+    dim_mults=(1, 2, 4, 8),
+    cond_on_text_encodings = True,
 ).cuda()
 
 unet2 = Unet(
     dim = 16,
     image_embed_dim = 512,
-    text_embed_dim = 512,
     cond_dim = 128,
     channels = 3,
     dim_mults = (1, 2, 4, 8, 16),
-    cond_on_text_encodings = True
 ).cuda()
 
 decoder = Decoder(
     unet = (unet1, unet2),
     image_sizes = (128, 256),
     clip = clip,
-    timesteps = 1000,
-    condition_on_text_encodings = True
+    timesteps = 1000
 ).cuda()
 
 decoder_trainer = DecoderTrainer(
@@ -903,8 +914,8 @@ for unet_number in (1, 2):
 # after much training
 # you can sample from the exponentially moving averaged unets as so
 
-mock_image_embed = torch.randn(4, 512).cuda()
-images = decoder_trainer.sample(mock_image_embed, text = text) # (4, 3, 256, 256)
+mock_image_embed = torch.randn(32, 512).cuda()
+images = decoder_trainer.sample(image_embed = mock_image_embed, text = text) # (4, 3, 256, 256)
 ```
 
 ### Diffusion Prior Training
@@ -1112,7 +1123,8 @@ For detailed information on training the diffusion prior, please refer to the [d
 - [x] allow for unet to be able to condition non-cross attention style as well
 - [x] speed up inference, read up on papers (ddim)
 - [x] add inpainting ability using resampler from repaint paper https://arxiv.org/abs/2201.09865
-- [ ] try out the nested unet from https://arxiv.org/abs/2005.09007 after hearing several positive testimonies from researchers, for segmentation anyhow
+- [x] add the final combination of upsample feature maps, used in unet squared, seems to have an effect in local experiments
+- [ ] consider elucidated dalle2 https://arxiv.org/abs/2206.00364
 - [ ] interface out the vqgan-vae so a pretrained one can be pulled off the shelf to validate latent diffusion + DALL-E2
 
 ## Citations
@@ -1241,4 +1253,15 @@ For detailed information on training the diffusion prior, please refer to the [d
 }
 ```
 
+```bibtex
+@misc{chen2022analog,
+    title   = {Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning},
+    author  = {Ting Chen and Ruixiang Zhang and Geoffrey Hinton},
+    year    = {2022},
+    eprint  = {2208.04202},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV}
+}
+```
+
 *Creating noise from data is easy; creating data from noise is generative modeling.* - <a href="https://arxiv.org/abs/2011.13456">Yang Song's paper</a>

dalle2_pytorch/trackers.py

@@ -528,8 +528,12 @@ class Tracker:
         elif save_type == 'model':
             if isinstance(trainer, DiffusionPriorTrainer):
                 prior = trainer.ema_diffusion_prior.ema_model if trainer.use_ema else trainer.diffusion_prior
-                state_dict = trainer.accelerator.unwrap_model(prior).state_dict()
-                torch.save(state_dict, file_path)
+                prior: DiffusionPrior = trainer.accelerator.unwrap_model(prior)
+                # Remove CLIP if it is part of the model
+                original_clip = prior.clip
+                prior.clip = None
+                model_state_dict = prior.state_dict()
+                prior.clip = original_clip
             elif isinstance(trainer, DecoderTrainer):
                 decoder: Decoder = trainer.accelerator.unwrap_model(trainer.decoder)
                 # Remove CLIP if it is part of the model

dalle2_pytorch/train_configs.py

@@ -1,7 +1,7 @@
 import json
 from torchvision import transforms as T
 from pydantic import BaseModel, validator, root_validator
-from typing import List, Iterable, Optional, Union, Tuple, Dict, Any
+from typing import List, Optional, Union, Tuple, Dict, Any, TypeVar
 
 from x_clip import CLIP as XCLIP
 from coca_pytorch import CoCa
@@ -25,11 +25,9 @@ def exists(val):
 def default(val, d):
     return val if exists(val) else d
 
-def ListOrTuple(inner_type):
-    return Union[List[inner_type], Tuple[inner_type]]
-
-def SingularOrIterable(inner_type):
-    return Union[inner_type, ListOrTuple(inner_type)]
+InnerType = TypeVar('InnerType')
+ListOrTuple = Union[List[InnerType], Tuple[InnerType]]
+SingularOrIterable = Union[InnerType, ListOrTuple[InnerType]]
 
 # general pydantic classes
 
@@ -222,13 +220,13 @@ class TrainDiffusionPriorConfig(BaseModel):
 
 class UnetConfig(BaseModel):
     dim: int
-    dim_mults: ListOrTuple(int)
+    dim_mults: ListOrTuple[int]
     image_embed_dim: int = None
     text_embed_dim: int = None
     cond_on_text_encodings: bool = None
     cond_dim: int = None
     channels: int = 3
-    self_attn: ListOrTuple(int)
+    self_attn: ListOrTuple[int]
     attn_dim_head: int = 32
     attn_heads: int = 16
     init_cross_embed: bool = True
@@ -237,16 +235,16 @@ class UnetConfig(BaseModel):
         extra = "allow"
 
 class DecoderConfig(BaseModel):
-    unets: ListOrTuple(UnetConfig)
+    unets: ListOrTuple[UnetConfig]
     image_size: int = None
-    image_sizes: ListOrTuple(int) = None
+    image_sizes: ListOrTuple[int] = None
     clip: Optional[AdapterConfig]   # The clip model to use if embeddings are not provided
     channels: int = 3
     timesteps: int = 1000
-    sample_timesteps: Optional[SingularOrIterable(int)] = None
+    sample_timesteps: Optional[SingularOrIterable[int]] = None
     loss_type: str = 'l2'
-    beta_schedule: ListOrTuple(str) = 'cosine'
-    learned_variance: bool = True
+    beta_schedule: ListOrTuple[str] = None  # None means all cosine
+    learned_variance: SingularOrIterable[bool] = True
     image_cond_drop_prob: float = 0.1
     text_cond_drop_prob: float = 0.5
 
@@ -305,11 +303,11 @@ class DecoderDataConfig(BaseModel):
 
 class DecoderTrainConfig(BaseModel):
     epochs: int = 20
-    lr: SingularOrIterable(float) = 1e-4
-    wd: SingularOrIterable(float) = 0.01
-    warmup_steps: Optional[SingularOrIterable(int)] = None
+    lr: SingularOrIterable[float] = 1e-4
+    wd: SingularOrIterable[float] = 0.01
+    warmup_steps: Optional[SingularOrIterable[int]] = None
     find_unused_parameters: bool = True
-    max_grad_norm: SingularOrIterable(float) = 0.5
+    max_grad_norm: SingularOrIterable[float] = 0.5
     save_every_n_samples: int = 100000
     n_sample_images: int = 6                       # The number of example images to produce when sampling the train and test dataset
     cond_scale: Union[float, List[float]] = 1.0
@@ -320,7 +318,7 @@ class DecoderTrainConfig(BaseModel):
     use_ema: bool = True
     ema_beta: float = 0.999
     amp: bool = False
-    unet_training_mask: ListOrTuple(bool) = None   # If None, use all unets
+    unet_training_mask: ListOrTuple[bool] = None   # If None, use all unets
 
 class DecoderEvaluateConfig(BaseModel):
     n_evaluation_samples: int = 1000

dalle2_pytorch/trainer.py

@@ -174,7 +174,7 @@ class DiffusionPriorTrainer(nn.Module):
     def __init__(
         self,
         diffusion_prior,
-        accelerator,
+        accelerator = None,
         use_ema = True,
         lr = 3e-4,
         wd = 1e-2,
@@ -186,8 +186,12 @@ class DiffusionPriorTrainer(nn.Module):
     ):
         super().__init__()
         assert isinstance(diffusion_prior, DiffusionPrior)
-        assert isinstance(accelerator, Accelerator)
+
         ema_kwargs, kwargs = groupby_prefix_and_trim('ema_', kwargs)
+        accelerator_kwargs, kwargs = groupby_prefix_and_trim('accelerator_', kwargs)
+
+        if not exists(accelerator):
+            accelerator = Accelerator(**accelerator_kwargs)
 
         # assign some helpful member vars
 
@@ -300,7 +304,7 @@ class DiffusionPriorTrainer(nn.Module):
 
         # all processes need to load checkpoint. no restriction here
         if isinstance(path_or_state, str):
-            path = Path(path)
+            path = Path(path_or_state)
             assert path.exists()
             loaded_obj = torch.load(str(path), map_location=self.device)
 

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.0.1'
+__version__ = '1.6.4'

setup.py

@@ -26,7 +26,7 @@ setup(
   install_requires=[
     'accelerate',
     'click',
-    'clip-anytorch',
+    'clip-anytorch>=2.4.0',
     'coca-pytorch>=0.0.5',
     'ema-pytorch>=0.0.7',
     'einops>=0.4',