backwards pass is not recommended under the autocast context, per pytorch docs

README.md

@@ -706,7 +706,7 @@ mock_image_embed = torch.randn(1, 512).cuda()
 images = decoder.sample(mock_image_embed) # (1, 3, 1024, 1024)
 ```
 
-## Training wrapper
+## Training wrapper (wip)
 
 ### Decoder Training
 
@@ -814,8 +814,8 @@ clip = CLIP(
 
 # mock data
 
-text = torch.randint(0, 49408, (32, 256)).cuda()
+text = torch.randint(0, 49408, (4, 256)).cuda()
-images = torch.randn(32, 3, 256, 256).cuda()
+images = torch.randn(4, 3, 256, 256).cuda()
 
 # prior networks (with transformer)
 
@@ -842,7 +842,7 @@ diffusion_prior_trainer = DiffusionPriorTrainer(
     ema_update_every = 10,
 )
 
-loss = diffusion_prior_trainer(text, images, max_batch_size = 4)
+loss = diffusion_prior_trainer(text, images)
 diffusion_prior_trainer.update()  # this will update the optimizer as well as the exponential moving averaged diffusion prior
 
 # after much of the above three lines in a loop
@@ -851,57 +851,6 @@ diffusion_prior_trainer.update()  # this will update the optimizer as well as th
 image_embeds = diffusion_prior_trainer.sample(text) # (4, 512) - exponential moving averaged image embeddings
 ```
 
-## Bonus
-
-### Unconditional Training
-
-The repository also contains the means to train unconditional DDPM model, or even cascading DDPMs. You simply have to set `unconditional = True` in the `Decoder`
-
-ex.
-
-```python
-import torch
-from dalle2_pytorch import Unet, Decoder
-
-# unet for the cascading ddpm
-
-unet1 = Unet(
-    dim = 128,
-    dim_mults=(1, 2, 4, 8)
-).cuda()
-
-unet2 = Unet(
-    dim = 32,
-    dim_mults = (1, 2, 4, 8, 16)
-).cuda()
-
-# decoder, which contains the unets
-
-decoder = Decoder(
-    unet = (unet1, unet2),
-    image_sizes = (256, 512),  # first unet up to 256px, then second to 512px
-    timesteps = 1000,
-    unconditional = True
-).cuda()
-
-# mock images (get a lot of this)
-
-images = torch.randn(1, 3, 512, 512).cuda()
-
-# feed images into decoder
-
-for i in (1, 2):
-    loss = decoder(images, unet_number = i)
-    loss.backward()
-
-# do the above for many many many many steps
-# then it will learn to generate images
-
-images = decoder.sample(batch_size = 2) # (2, 3, 512, 512)
-```
-
-## Dataloaders
-
 ### Decoder Dataloaders
 
 In order to make loading data simple and efficient, we include some general dataloaders that can be used to train portions of the network.
@@ -946,14 +895,14 @@ dataset = ImageEmbeddingDataset(
 )
 ```
 
-### Scripts (wip)
+## Scripts
 
-#### `train_diffusion_prior.py`
+### Using the `train_diffusion_prior.py` script
 
 This script allows training the DiffusionPrior on pre-computed text and image embeddings. The working example below elucidates this process.
 Please note that the script internally passes text_embed and image_embed to the DiffusionPrior, unlike the example below.
 
-#### Usage
+### Usage 
 
 ```bash
 $ python train_diffusion_prior.py
@@ -961,49 +910,58 @@ $ python train_diffusion_prior.py
 
 The most significant parameters for the script are as follows:
 
-- `image-embed-url`, default = `"https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/img_emb/"`
+--image-embed-url, default = "https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/img_emb/")
 
-- `text-embed-url`, default = `"https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/text_emb/"`
+--text-embed-url, default = "https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/text_emb/")
 
-- `image-embed-dim`, default = `768` - 768 corresponds to the ViT iL/14 embedding size,change it to what your chosen ViT generates
+--image-embed-dim, default=768 - 768 corresponds to the ViT iL/14 embedding size,change it to what your chosen ViT generates
 
-- `learning-rate`, default = `1.1e-4`
+--learning-rate, default=1.1e-4
 
-- `weight-decay`,  default = `6.02e-2`
+--weight-decay,  default=6.02e-2
 
-- `max-grad-norm`, default = `0.5`
+--max-grad-norm, default=0.5
 
-- `batch-size`, default = `10 ** 4`
+--batch-size, default=10 ** 4
 
-- `num-epochs`, default = `5`
+--num-epochs, default=5
 
-- `clip`, default = `None` # Signals the prior to use pre-computed embeddings
+--clip, default=None # Signals the prior to use pre-computed embeddings
 
-#### Loading and Saving the DiffusionPrior model
+### Sample wandb run log
+
+Please find a sample wandb run log at : https://wandb.ai/laion/diffusion-prior/runs/1blxu24j
+
+### Loading and saving the Diffusion Prior model
 
 Two methods are provided, load_diffusion_model and save_diffusion_model, the names being self-explanatory. 
 
-```python
+## from dalle2_pytorch.train import load_diffusion_model, save_diffusion_model
-from dalle2_pytorch.train import load_diffusion_model, save_diffusion_model
-```
-
-##### Loading
 
     load_diffusion_model(dprior_path, device) 
+
         dprior_path : path to saved model(.pth)
+    
         device      : the cuda device you're running on
     
-##### Saving
-
     save_diffusion_model(save_path, model, optimizer, scaler, config, image_embed_dim)
+    
         save_path : path to save at
+    
         model     : object of Diffusion_Prior
+    
         optimizer : optimizer object - see train_diffusion_prior.py for how to create one. 
+    
             e.g: optimizer = get_optimizer(diffusion_prior.net.parameters(), wd=weight_decay, lr=learning_rate)
+    
         scaler    : a GradScaler object.
+    
             e.g: scaler = GradScaler(enabled=amp)
+    
         config    : config object created in train_diffusion_prior.py - see file for example. 
+    
         image_embed_dim - the dimension of the image_embedding
+    
             e.g: 768
 
 ## CLI (wip)
@@ -1049,7 +1007,6 @@ Once built, images will be saved to the same directory the command is invoked
 - [x] make sure the cascading ddpm in the repository can be trained unconditionally, offer a one-line CLI tool for training on a folder of images
 - [x] bring in cross-scale embedding from iclr paper https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/crossformer.py#L14
 - [x] cross embed layers for downsampling, as an option
-- [x] use an experimental tracker agnostic setup, as done <a href="https://github.com/lucidrains/tf-bind-transformer#simple-trainer-class-for-fine-tuning">here</a>
 - [ ] become an expert with unets, cleanup unet code, make it fully configurable, port all learnings over to https://github.com/lucidrains/x-unet (test out unet² in ddpm repo) - consider https://github.com/lucidrains/uformer-pytorch attention-based unet
 - [ ] transcribe code to Jax, which lowers the activation energy for distributed training, given access to TPUs
 - [ ] train on a toy task, offer in colab
@@ -1057,15 +1014,13 @@ Once built, images will be saved to the same directory the command is invoked
 - [ ] extend diffusion head to use diffusion-gan (potentially using lightweight-gan) to speed up inference
 - [ ] figure out if possible to augment with external memory, as described in https://arxiv.org/abs/2204.11824
 - [ ] test out grid attention in cascading ddpm locally, decide whether to keep or remove
+- [ ] use an experimental tracker agnostic setup, as done <a href="https://github.com/lucidrains/tf-bind-transformer#simple-trainer-class-for-fine-tuning">here</a>
 - [ ] interface out the vqgan-vae so a pretrained one can be pulled off the shelf to validate latent diffusion + DALL-E2
 - [ ] make sure FILIP works with DALL-E2 from x-clip https://arxiv.org/abs/2111.07783
 - [ ] offer save / load methods on the trainer classes to automatically take care of state dicts for scalers / optimizers / saving versions and checking for breaking changes
 - [ ] bring in skip-layer excitatons (from lightweight gan paper) to see if it helps for either decoder of unet or vqgan-vae training
 - [ ] decoder needs one day worth of refactor for tech debt
 - [ ] allow for unet to be able to condition non-cross attention style as well
-- [ ] for all model classes with hyperparameters that changes the network architecture, make it requirement that they must expose a config property, and write a simple function that asserts that it restores the object correctly
-- [ ] for both diffusion prior and decoder, all exponential moving averaged models needs to be saved and restored as well (as well as the step number)
-- [ ] read the paper, figure it out, and build it https://github.com/lucidrains/DALLE2-pytorch/issues/89
 
 ## Citations
 
@@ -1154,13 +1109,4 @@ Once built, images will be saved to the same directory the command is invoked
 }
 ```
 
-```bibtex
-@article{ho2021cascaded,
-    title   = {Cascaded Diffusion Models for High Fidelity Image Generation},
-    author  = {Ho, Jonathan and Saharia, Chitwan and Chan, William and Fleet, David J and Norouzi, Mohammad and Salimans, Tim},
-    journal = {arXiv preprint arXiv:2106.15282},
-    year    = {2021}
-}
-```
-
 *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/__init__.py

@@ -1,6 +1,6 @@
 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
+from dalle2_pytorch.train import DecoderTrainer, DiffusionPriorTrainer
 
 from dalle2_pytorch.vqgan_vae import VQGanVAE
 from x_clip import CLIP

dalle2_pytorch/dalle2_pytorch.py

@@ -61,9 +61,6 @@ def default(val, d):
 def cast_tuple(val, length = 1):
     return val if isinstance(val, tuple) else ((val,) * length)
 
-def module_device(module):
-    return next(module.parameters()).device
-
 @contextmanager
 def null_context(*args, **kwargs):
     yield
@@ -797,7 +794,7 @@ class DiffusionPriorNetwork(nn.Module):
         text_embed,
         text_encodings = None,
         mask = None,
-        cond_drop_prob = 0.
+        cond_drop_prob = 0.2
     ):
         batch, dim, device, dtype = *image_embed.shape, image_embed.device, image_embed.dtype
 
@@ -904,7 +901,6 @@ class DiffusionPrior(BaseGaussianDiffusion):
         self.channels = default(image_channels, lambda: clip.image_channels)
 
         self.cond_drop_prob = cond_drop_prob
-        self.can_classifier_guidance = cond_drop_prob > 0.
         self.condition_on_text_encodings = condition_on_text_encodings
 
         # in paper, they do not predict the noise, but predict x0 directly for image embedding, claiming empirically better results. I'll just offer both.
@@ -918,10 +914,8 @@ class DiffusionPrior(BaseGaussianDiffusion):
         self.training_clamp_l2norm = training_clamp_l2norm
         self.init_image_embed_l2norm = init_image_embed_l2norm
 
-    def p_mean_variance(self, x, t, text_cond, clip_denoised = False, cond_scale = 1.):
+    def p_mean_variance(self, x, t, text_cond, clip_denoised: bool):
-        assert not (cond_scale != 1. and not self.can_classifier_guidance), 'the model was not trained with conditional dropout, and thus one cannot use classifier free guidance (cond_scale anything other than 1)'
+        pred = self.net(x, t, **text_cond)
-
-        pred = self.net.forward_with_cond_scale(x, t, cond_scale = cond_scale, **text_cond)
 
         if self.predict_x_start:
             x_recon = pred
@@ -939,17 +933,17 @@ class DiffusionPrior(BaseGaussianDiffusion):
         model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
         return model_mean, posterior_variance, posterior_log_variance
 
-    @torch.no_grad()
+    @torch.inference_mode()
-    def p_sample(self, x, t, text_cond = None, clip_denoised = True, repeat_noise = False, cond_scale = 1.):
+    def p_sample(self, x, t, text_cond = None, clip_denoised = True, repeat_noise = False):
         b, *_, device = *x.shape, x.device
-        model_mean, _, model_log_variance = self.p_mean_variance(x = x, t = t, text_cond = text_cond, clip_denoised = clip_denoised, cond_scale = cond_scale)
+        model_mean, _, model_log_variance = self.p_mean_variance(x = x, t = t, text_cond = text_cond, clip_denoised = clip_denoised)
         noise = noise_like(x.shape, device, repeat_noise)
         # no noise when t == 0
         nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
         return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
 
-    @torch.no_grad()
+    @torch.inference_mode()
-    def p_sample_loop(self, shape, text_cond, cond_scale = 1.):
+    def p_sample_loop(self, shape, text_cond):
         device = self.betas.device
 
         b = shape[0]
@@ -960,7 +954,7 @@ class DiffusionPrior(BaseGaussianDiffusion):
 
         for i in tqdm(reversed(range(0, self.num_timesteps)), desc='sampling loop time step', total=self.num_timesteps):
             times = torch.full((b,), i, device = device, dtype = torch.long)
-            image_embed = self.p_sample(image_embed, times, text_cond = text_cond, cond_scale = cond_scale)
+            image_embed = self.p_sample(image_embed, times, text_cond = text_cond)
 
         return image_embed
 
@@ -984,21 +978,21 @@ class DiffusionPrior(BaseGaussianDiffusion):
         loss = self.loss_fn(pred, target)
         return loss
 
-    @torch.no_grad()
+    @torch.inference_mode()
     @eval_decorator
-    def sample_batch_size(self, batch_size, text_cond, cond_scale = 1.):
+    def sample_batch_size(self, batch_size, text_cond):
         device = self.betas.device
         shape = (batch_size, self.image_embed_dim)
 
         img = torch.randn(shape, device = device)
 
         for i in tqdm(reversed(range(0, self.num_timesteps)), desc = 'sampling loop time step', total = self.num_timesteps):
-            img = self.p_sample(img, torch.full((batch_size,), i, device = device, dtype = torch.long), text_cond = text_cond, cond_scale = cond_scale)
+            img = self.p_sample(img, torch.full((batch_size,), i, device = device, dtype = torch.long), text_cond = text_cond)
         return img
 
-    @torch.no_grad()
+    @torch.inference_mode()
     @eval_decorator
-    def sample(self, text, num_samples_per_batch = 2, cond_scale = 1.):
+    def sample(self, text, num_samples_per_batch = 2):
         # in the paper, what they did was
         # sample 2 image embeddings, choose the top 1 similarity, as judged by CLIP
         text = repeat(text, 'b ... -> (b r) ...', r = num_samples_per_batch)
@@ -1013,7 +1007,7 @@ class DiffusionPrior(BaseGaussianDiffusion):
         if self.condition_on_text_encodings:
             text_cond = {**text_cond, 'text_encodings': text_encodings, 'mask': text_mask}
 
-        image_embeds = self.p_sample_loop((batch_size, image_embed_dim), text_cond = text_cond, cond_scale = cond_scale)
+        image_embeds = self.p_sample_loop((batch_size, image_embed_dim), text_cond = text_cond)
 
         # retrieve original unscaled image embed
 
@@ -1311,7 +1305,7 @@ class Unet(nn.Module):
         self,
         dim,
         *,
-        image_embed_dim = None,
+        image_embed_dim,
         text_embed_dim = None,
         cond_dim = None,
         num_image_tokens = 4,
@@ -1383,7 +1377,7 @@ class Unet(nn.Module):
         self.image_to_cond = nn.Sequential(
             nn.Linear(image_embed_dim, cond_dim * num_image_tokens),
             Rearrange('b (n d) -> b n d', n = num_image_tokens)
-        ) if cond_on_image_embeds and image_embed_dim != cond_dim else nn.Identity()
+        ) if image_embed_dim != cond_dim else nn.Identity()
 
         self.norm_cond = nn.LayerNorm(cond_dim)
         self.norm_mid_cond = nn.LayerNorm(cond_dim)
@@ -1393,8 +1387,7 @@ class Unet(nn.Module):
         self.text_to_cond = None
 
         if cond_on_text_encodings:
-            assert exists(text_embed_dim), 'text_embed_dim must be given to the unet if cond_on_text_encodings is True'
+            self.text_to_cond = nn.LazyLinear(cond_dim) if not exists(text_embed_dim) else nn.Linear(text_embed_dim, cond_dim)
-            self.text_to_cond = nn.Linear(text_embed_dim, cond_dim)
 
         # finer control over whether to condition on image embeddings and text encodings
         # so one can have the latter unets in the cascading DDPMs only focus on super-resoluting
@@ -1708,7 +1701,7 @@ class Decoder(BaseGaussianDiffusion):
         self.unconditional = unconditional
         assert not (condition_on_text_encodings and unconditional), 'unconditional decoder image generation cannot be set to True if conditioning on text is present'
 
-        assert self.unconditional or (exists(clip) ^ exists(image_size)), 'either CLIP is supplied, or you must give the image_size and channels (usually 3 for RGB)'
+        assert exists(clip) ^ exists(image_size), 'either CLIP is supplied, or you must give the image_size and channels (usually 3 for RGB)'
 
         self.clip = None
         if exists(clip):
@@ -1799,7 +1792,6 @@ class Decoder(BaseGaussianDiffusion):
 
         self.image_cond_drop_prob = image_cond_drop_prob
         self.text_cond_drop_prob = text_cond_drop_prob
-        self.can_classifier_guidance = image_cond_drop_prob > 0. or text_cond_drop_prob > 0.
 
         # whether to clip when sampling
 
@@ -1819,19 +1811,13 @@ class Decoder(BaseGaussianDiffusion):
             unet = self.get_unet(unet_number)
 
         self.cuda()
-
-        devices = [module_device(unet) for unet in self.unets]
         self.unets.cpu()
+
         unet.cuda()
-
         yield
-
+        unet.cpu()
-        for unet, device in zip(self.unets, devices):
-            unet.to(device)
 
     def p_mean_variance(self, unet, x, t, image_embed, text_encodings = None, text_mask = None, lowres_cond_img = None, clip_denoised = True, predict_x_start = False, learned_variance = False, cond_scale = 1., model_output = None):
-        assert not (cond_scale != 1. and not self.can_classifier_guidance), 'the decoder was not trained with conditional dropout, and thus one cannot use classifier free guidance (cond_scale anything other than 1)'
-
         pred = default(model_output, lambda: unet.forward_with_cond_scale(x, t, image_embed = image_embed, text_encodings = text_encodings, text_mask = text_mask, cond_scale = cond_scale, lowres_cond_img = lowres_cond_img))
 
         if learned_variance:
@@ -1860,7 +1846,7 @@ class Decoder(BaseGaussianDiffusion):
 
         return model_mean, posterior_variance, posterior_log_variance
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def p_sample(self, unet, x, t, image_embed, text_encodings = None, text_mask = None, cond_scale = 1., lowres_cond_img = None, predict_x_start = False, learned_variance = False, clip_denoised = True, repeat_noise = False):
         b, *_, device = *x.shape, x.device
         model_mean, _, model_log_variance = self.p_mean_variance(unet, x = x, t = t, image_embed = image_embed, text_encodings = text_encodings, text_mask = text_mask, cond_scale = cond_scale, lowres_cond_img = lowres_cond_img, clip_denoised = clip_denoised, predict_x_start = predict_x_start, learned_variance = learned_variance)
@@ -1869,7 +1855,7 @@ class Decoder(BaseGaussianDiffusion):
         nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
         return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def p_sample_loop(self, unet, shape, image_embed, predict_x_start = False, learned_variance = False, clip_denoised = True, lowres_cond_img = None, text_encodings = None, text_mask = None, cond_scale = 1):
         device = self.betas.device
 
@@ -1962,14 +1948,12 @@ class Decoder(BaseGaussianDiffusion):
 
         return loss + vb_loss
 
-    @torch.no_grad()
+    @torch.inference_mode()
     @eval_decorator
     def sample(
         self,
         image_embed = None,
         text = None,
-        text_mask = None,
-        text_encodings = None,
         batch_size = 1,
         cond_scale = 1.,
         stop_at_unet_number = None
@@ -1979,8 +1963,8 @@ class Decoder(BaseGaussianDiffusion):
         if not self.unconditional:
             batch_size = image_embed.shape[0]
 
-        if exists(text) and not exists(text_encodings) and not self.unconditional:
+        text_encodings = text_mask = None
-            assert exist(self.clip)
+        if exists(text):
             _, text_encodings, text_mask = self.clip.embed_text(text)
 
         assert not (self.condition_on_text_encodings and not exists(text_encodings)), 'text or text encodings must be passed into decoder if specified'
@@ -2004,7 +1988,8 @@ class Decoder(BaseGaussianDiffusion):
                 image_size = vae.get_encoded_fmap_size(image_size)
                 shape = (batch_size, vae.encoded_dim, image_size, image_size)
 
-                lowres_cond_img = maybe(vae.encode)(lowres_cond_img)
+                if exists(lowres_cond_img):
+                    lowres_cond_img = vae.encode(lowres_cond_img)
 
                 img = self.p_sample_loop(
                     unet,
@@ -2032,7 +2017,6 @@ class Decoder(BaseGaussianDiffusion):
         text = None,
         image_embed = None,
         text_encodings = None,
-        text_mask = None,
         unet_number = None
     ):
         assert not (len(self.unets) > 1 and not exists(unet_number)), f'you must specify which unet you want trained, from a range of 1 to {len(self.unets)}, if you are training cascading DDPM (multiple unets)'
@@ -2053,11 +2037,12 @@ class Decoder(BaseGaussianDiffusion):
 
         times = torch.randint(0, self.num_timesteps, (b,), device = device, dtype = torch.long)
 
-        if not exists(image_embed) and not self.unconditional:
+        if not exists(image_embed):
             assert exists(self.clip), 'if you want to derive CLIP image embeddings automatically, you must supply `clip` to the decoder on init'
             image_embed, _ = self.clip.embed_image(image)
 
-        if exists(text) and not exists(text_encodings) and not self.unconditional:
+        text_encodings = text_mask = None
+        if exists(text) and not exists(text_encodings):
             assert exists(self.clip), 'if you are passing in raw text, you need to supply `clip` to the decoder'
             _, text_encodings, text_mask = self.clip.embed_text(text)
 
@@ -2078,7 +2063,9 @@ class Decoder(BaseGaussianDiffusion):
         vae.eval()
         with torch.no_grad():
             image = vae.encode(image)
-            lowres_cond_img = maybe(vae.encode)(lowres_cond_img)
+
+            if exists(lowres_cond_img):
+                lowres_cond_img = vae.encode(lowres_cond_img)
 
         return self.p_losses(unet, image, times, image_embed = image_embed, text_encodings = text_encodings, text_mask = text_mask, lowres_cond_img = lowres_cond_img, predict_x_start = predict_x_start, learned_variance = learned_variance)
 
@@ -2103,23 +2090,22 @@ class DALLE2(nn.Module):
 
         self.to_pil = T.ToPILImage()
 
-    @torch.no_grad()
+    @torch.inference_mode()
     @eval_decorator
     def forward(
         self,
         text,
         cond_scale = 1.,
-        prior_cond_scale = 1.,
         return_pil_images = False
     ):
-        device = module_device(self)
+        device = next(self.parameters()).device
         one_text = isinstance(text, str) or (not is_list_str(text) and text.shape[0] == 1)
 
         if isinstance(text, str) or is_list_str(text):
             text = [text] if not isinstance(text, (list, tuple)) else text
             text = tokenizer.tokenize(text).to(device)
 
-        image_embed = self.prior.sample(text, num_samples_per_batch = self.prior_num_samples, cond_scale = prior_cond_scale)
+        image_embed = self.prior.sample(text, num_samples_per_batch = self.prior_num_samples)
 
         text_cond = text if self.decoder_need_text_cond else None
         images = self.decoder.sample(image_embed, text = text_cond, cond_scale = cond_scale)

dalle2_pytorch/dataloaders/__init__.py

@@ -1,2 +1 @@
-from dalle2_pytorch.dataloaders.decoder_loader import ImageEmbeddingDataset, create_image_embedding_dataloader
+from dalle2_pytorch.dataloaders.decoder_loader import ImageEmbeddingDataset, create_image_embedding_dataloader
-from dalle2_pytorch.dataloaders.embedding_wrapper import make_splits

dalle2_pytorch/dataloaders/embedding_wrapper.py

@@ -1,180 +0,0 @@
-from torch.utils.data import IterableDataset
-from torch import from_numpy
-from clip import tokenize
-from embedding_reader import EmbeddingReader
-
-
-class PriorEmbeddingLoader(IterableDataset):
-    def __init__(
-        self,
-        text_conditioned: bool,
-        batch_size: int,
-        start: int,
-        stop: int,
-        image_reader,
-        text_reader: EmbeddingReader = None,
-        device: str = "cpu",
-    ) -> None:
-        super(PriorEmbeddingLoader).__init__()
-
-        self.text_conditioned = text_conditioned
-
-        if not self.text_conditioned:
-            self.text_reader = text_reader
-
-        self.image_reader = image_reader
-        self.batch_size = batch_size
-        self.start = start
-        self.stop = stop
-        self.device = device
-
-    def __iter__(self):
-        self.n = 0
-        loader_args = dict(
-            batch_size=self.batch_size,
-            start=self.start,
-            end=self.stop,
-            show_progress=False,
-        )
-        if self.text_conditioned:
-            self.loader = self.image_reader(**loader_args)
-        else:
-            self.loader = zip(
-                self.image_reader(**loader_args), self.text_reader(**loader_args)
-            )
-        return self
-
-    def __next__(self):
-        try:
-            return self.get_sample()
-        except StopIteration:
-            raise StopIteration
-
-    def get_sample(self):
-        """
-        pre-proocess data from either reader into a common format
-        """
-        self.n += 1
-
-        if self.text_conditioned:
-            image_embedding, caption = next(self.loader)
-
-            image_embedding = from_numpy(image_embedding).to(self.device)
-            tokenized_caption = tokenize(
-                caption["caption"].to_list(), truncate=True
-            ).to(self.device)
-
-            return image_embedding, tokenized_caption
-
-        else:
-            (image_embedding, _), (text_embedding, _) = next(self.loader)
-
-            image_embedding = from_numpy(image_embedding).to(self.device)
-            text_embedding = from_numpy(text_embedding).to(self.device)
-
-            return image_embedding, text_embedding
-
-
-def make_splits(
-    text_conditioned: bool,
-    batch_size: int,
-    num_data_points: int,
-    train_split: float,
-    eval_split: float,
-    device: str,
-    img_url: str,
-    meta_url: str = None,
-    txt_url: str = None,
-):
-
-    assert img_url is not None, "Must supply some image embeddings"
-
-    if text_conditioned:
-        assert meta_url is not None, "Must supply metadata url if text-conditioning"
-        image_reader = EmbeddingReader(
-            embeddings_folder=img_url,
-            file_format="parquet_npy",
-            meta_columns=["caption"],
-            metadata_folder=meta_url,
-        )
-
-        # compute split points
-        if num_data_points > image_reader.count:
-            print("Specified point count is larger than the number of points available...defaulting to max length of reader.")
-            num_data_points = image_reader.count
-
-        train_set_size = int(train_split * num_data_points)
-        eval_set_size = int(eval_split * num_data_points)
-        eval_stop = int(train_set_size + eval_set_size)
-
-        train_loader = PriorEmbeddingLoader(
-            text_conditioned=text_conditioned,
-            image_reader=image_reader,
-            batch_size=batch_size,
-            start=0,
-            stop=train_set_size,
-            device=device,
-        )
-        eval_loader = PriorEmbeddingLoader(
-            text_conditioned=text_conditioned,
-            image_reader=image_reader,
-            batch_size=batch_size,
-            start=train_set_size,
-            stop=eval_stop,
-            device=device,
-        )
-        test_loader = PriorEmbeddingLoader(
-            text_conditioned=text_conditioned,
-            image_reader=image_reader,
-            batch_size=batch_size,
-            start=eval_stop,
-            stop=int(num_data_points),
-            device=device,
-        )
-
-    else:
-        assert (
-            txt_url is not None
-        ), "Must supply text embedding url if not text-conditioning"
-
-        image_reader = EmbeddingReader(img_url, file_format="npy")
-        text_reader = EmbeddingReader(txt_url, file_format="npy")
-
-        # compute split points
-        if num_data_points > image_reader.count:
-            print("Specified point count is larger than the number of points available...defaulting to max length of reader.")
-            num_data_points = image_reader.count
-
-        train_set_size = int(train_split * num_data_points)
-        eval_set_size = int(eval_split * num_data_points)
-        eval_stop = int(train_set_size + eval_set_size)
-
-        train_loader = PriorEmbeddingLoader(
-            text_conditioned=text_conditioned,
-            image_reader=image_reader,
-            text_reader=text_reader,
-            batch_size=batch_size,
-            start=0,
-            stop=train_set_size,
-            device=device,
-        )
-        eval_loader = PriorEmbeddingLoader(
-            text_conditioned=text_conditioned,
-            image_reader=image_reader,
-            text_reader=text_reader,
-            batch_size=batch_size,
-            start=train_set_size,
-            stop=eval_stop,
-            device=device,
-        )
-        test_loader = PriorEmbeddingLoader(
-            text_conditioned=text_conditioned,
-            image_reader=image_reader,
-            text_reader=text_reader,
-            batch_size=batch_size,
-            start=eval_stop,
-            stop=int(num_data_points),
-            device=device,
-        )
-
-    return train_loader, eval_loader, test_loader

dalle2_pytorch/trackers.py

@@ -1,49 +0,0 @@
-import os
-import torch
-from torch import nn
-
-# helper functions
-
-def exists(val):
-    return val is not None
-
-# base class
-
-class BaseTracker(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def init(self, config, **kwargs):
-        raise NotImplementedError
-
-    def log(self, log, **kwargs):
-        raise NotImplementedError
-
-# basic stdout class
-
-class ConsoleTracker(BaseTracker):
-    def init(self, **config):
-        print(config)
-
-    def log(self, log, **kwargs):
-        print(log)
-
-# basic wandb class
-
-class WandbTracker(BaseTracker):
-    def __init__(self):
-        super().__init__()
-        try:
-            import wandb
-        except ImportError as e:
-            print('`pip install wandb` to use the wandb experiment tracker')
-            raise e
-
-        os.environ["WANDB_SILENT"] = "true"
-        self.wandb = wandb
-
-    def init(self, **config):
-        self.wandb.init(**config)
-
-    def log(self, log, **kwargs):
-        self.wandb.log(log, **kwargs)

dalle2_pytorch/train.py

@@ -1,7 +1,7 @@
 import time
 import copy
 from math import ceil
-from functools import partial, wraps
+from functools import partial
 from collections.abc import Iterable
 
 import torch
@@ -11,8 +11,6 @@ from torch.cuda.amp import autocast, GradScaler
 from dalle2_pytorch.dalle2_pytorch import Decoder, DiffusionPrior
 from dalle2_pytorch.optimizer import get_optimizer
 
-import numpy as np
-
 # helper functions
 
 def exists(val):
@@ -47,29 +45,6 @@ def groupby_prefix_and_trim(prefix, d):
     kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items())))
     return kwargs_without_prefix, kwargs
 
-# decorators
-
-def cast_torch_tensor(fn):
-    @wraps(fn)
-    def inner(model, *args, **kwargs):
-        device = kwargs.pop('_device', next(model.parameters()).device)
-        cast_device = kwargs.pop('_cast_device', True)
-
-        kwargs_keys = kwargs.keys()
-        all_args = (*args, *kwargs.values())
-        split_kwargs_index = len(all_args) - len(kwargs_keys)
-        all_args = tuple(map(lambda t: torch.from_numpy(t) if exists(t) and isinstance(t, np.ndarray) else t, all_args))
-
-        if cast_device:
-            all_args = tuple(map(lambda t: t.to(device) if exists(t) and isinstance(t, torch.Tensor) else t, all_args))
-
-        args, kwargs_values = all_args[:split_kwargs_index], all_args[split_kwargs_index:]
-        kwargs = dict(tuple(zip(kwargs_keys, kwargs_values)))
-
-        out = fn(model, *args, **kwargs)
-        return out
-    return inner
-
 # gradient accumulation functions
 
 def split_iterable(it, split_size):
@@ -91,27 +66,18 @@ def split(t, split_size = None):
 
     return TypeError
 
-def find_first(cond, arr):
+def split_args_and_kwargs(x, *args, split_size = None, **kwargs):
-    for el in arr:
+    batch_size = len(x)
-        if cond(el):
-            return el
-    return None
-
-def split_args_and_kwargs(*args, split_size = None, **kwargs):
-    all_args = (*args, *kwargs.values())
-    len_all_args = len(all_args)
-    first_tensor = find_first(lambda t: isinstance(t, torch.Tensor), all_args)
-    assert exists(first_tensor)
-
-    batch_size = len(first_tensor)
     split_size = default(split_size, batch_size)
-    num_chunks = ceil(batch_size / split_size)
+    chunk_size = ceil(batch_size / split_size)
 
     dict_len = len(kwargs)
     dict_keys = kwargs.keys()
+    all_args = (x, *args, *kwargs.values())
+    len_all_args = len(all_args)
     split_kwargs_index = len_all_args - dict_len
 
-    split_all_args = [split(arg, split_size = split_size) if exists(arg) and isinstance(arg, (torch.Tensor, Iterable)) else ((arg,) * num_chunks) for arg in all_args]
+    split_all_args = [split(arg, split_size = split_size) if exists(arg) and isinstance(arg, (torch.Tensor, Iterable)) else ((arg,) * chunk_size) for arg in all_args]
     chunk_sizes = tuple(map(len, split_all_args[0]))
 
     for (chunk_size, *chunked_all_args) in tuple(zip(chunk_sizes, *split_all_args)):
@@ -151,7 +117,7 @@ def load_diffusion_model(dprior_path, device):
     # Load state dict from saved model
     diffusion_prior.load_state_dict(loaded_obj['model'])
 
-    return diffusion_prior, loaded_obj
+    return diffusion_prior
 
 def save_diffusion_model(save_path, model, optimizer, scaler, config, image_embed_dim):
     # Saving State Dict
@@ -262,8 +228,6 @@ class DiffusionPriorTrainer(nn.Module):
 
         self.max_grad_norm = max_grad_norm
 
-        self.register_buffer('step', torch.tensor([0.]))
-
     def update(self):
         if exists(self.max_grad_norm):
             self.scaler.unscale_(self.optimizer)
@@ -276,40 +240,34 @@ class DiffusionPriorTrainer(nn.Module):
         if self.use_ema:
             self.ema_diffusion_prior.update()
 
-        self.step += 1
+    @torch.inference_mode()
-
-    @torch.no_grad()
-    @cast_torch_tensor
     def p_sample_loop(self, *args, **kwargs):
         return self.ema_diffusion_prior.ema_model.p_sample_loop(*args, **kwargs)
 
-    @torch.no_grad()
+    @torch.inference_mode()
-    @cast_torch_tensor
     def sample(self, *args, **kwargs):
         return self.ema_diffusion_prior.ema_model.sample(*args, **kwargs)
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def sample_batch_size(self, *args, **kwargs):
         return self.ema_diffusion_prior.ema_model.sample_batch_size(*args, **kwargs)
 
-    @cast_torch_tensor
     def forward(
         self,
+        x,
         *args,
         max_batch_size = None,
         **kwargs
     ):
         total_loss = 0.
 
-        for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, split_size = max_batch_size, **kwargs):
+        for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(x, *args, split_size = max_batch_size, **kwargs):
             with autocast(enabled = self.amp):
                 loss = self.diffusion_prior(*chunked_args, **chunked_kwargs)
-                loss = loss * chunk_size_frac
 
+            loss = loss * chunk_size_frac
             total_loss += loss.item()
-
+            self.scaler.scale(loss).backward()
-            if self.training:
-                self.scaler.scale(loss).backward()
 
         return total_loss
 
@@ -335,6 +293,11 @@ class DecoderTrainer(nn.Module):
         self.num_unets = len(self.decoder.unets)
 
         self.use_ema = use_ema
+
+        if use_ema:
+            has_lazy_linear = any([type(module) == nn.LazyLinear for module in decoder.modules()])
+            assert not has_lazy_linear, 'you must set the text_embed_dim on your u-nets if you plan on doing automatic exponential moving average'
+
         self.ema_unets = nn.ModuleList([])
 
         self.amp = amp
@@ -365,8 +328,6 @@ class DecoderTrainer(nn.Module):
 
         self.max_grad_norm = max_grad_norm
 
-        self.register_buffer('step', torch.tensor([0.]))
-
     @property
     def unets(self):
         return nn.ModuleList([ema.ema_model for ema in self.ema_unets])
@@ -377,11 +338,8 @@ class DecoderTrainer(nn.Module):
         scaler = getattr(self, f'scaler{index}')
         return scaler.scale(loss)
 
-    def update(self, unet_number = None):
+    def update(self, unet_number):
-        if self.num_unets == 1:
+        assert 1 <= unet_number <= self.num_unets
-            unet_number = default(unet_number, 1)
-
-        assert exists(unet_number) and 1 <= unet_number <= self.num_unets
         index = unet_number - 1
         unet = self.decoder.unets[index]
 
@@ -400,10 +358,7 @@ class DecoderTrainer(nn.Module):
             ema_unet = self.ema_unets[index]
             ema_unet.update()
 
-        self.step += 1
-
     @torch.no_grad()
-    @cast_torch_tensor
     def sample(self, *args, **kwargs):
         if self.use_ema:
             trainable_unets = self.decoder.unets
@@ -420,27 +375,22 @@ class DecoderTrainer(nn.Module):
 
         return output
 
-    @cast_torch_tensor
     def forward(
         self,
-        *args,
+        x,
-        unet_number = None,
+        *,
+        unet_number,
         max_batch_size = None,
         **kwargs
     ):
-        if self.num_unets == 1:
-            unet_number = default(unet_number, 1)
-
         total_loss = 0.
 
-        for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, split_size = max_batch_size, **kwargs):
+        for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(x, split_size = max_batch_size, **kwargs):
             with autocast(enabled = self.amp):
                 loss = self.decoder(*chunked_args, unet_number = unet_number, **chunked_kwargs)
-                loss = loss * chunk_size_frac
 
+            loss = loss * chunk_size_frac
             total_loss += loss.item()
-
+            self.scale(loss, unet_number = unet_number).backward()
-            if self.training:
-                self.scale(loss, unet_number = unet_number).backward()
 
         return total_loss

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.2.41',
+  version = '0.2.28',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',
@@ -30,7 +30,6 @@ setup(
     'einops-exts>=0.0.3',
     'embedding-reader',
     'kornia>=0.5.4',
-    'numpy',
     'pillow',
     'resize-right>=0.0.2',
     'rotary-embedding-torch',

train_diffusion_prior.py

@@ -1,375 +1,354 @@
-from pathlib import Path
+import os
-import click
 import math
-import time
+import argparse
 import numpy as np
 
 import torch
-import clip
 from torch import nn
-
-from dalle2_pytorch.dataloaders import make_splits
-from dalle2_pytorch import DiffusionPrior, DiffusionPriorNetwork, OpenAIClipAdapter
-from dalle2_pytorch.trainer import DiffusionPriorTrainer, load_diffusion_model, save_diffusion_model, print_ribbon
-
-from dalle2_pytorch.trackers import ConsoleTracker, WandbTracker
-
 from embedding_reader import EmbeddingReader
+from dalle2_pytorch import DiffusionPrior, DiffusionPriorNetwork
+from dalle2_pytorch.train import load_diffusion_model, save_diffusion_model, print_ribbon
+from dalle2_pytorch.optimizer import get_optimizer
+from torch.cuda.amp import autocast,GradScaler
 
+import time
 from tqdm import tqdm
 
-# constants
+import wandb
+os.environ["WANDB_SILENT"] = "true"
+NUM_TEST_EMBEDDINGS = 100 # for cosine similarity reporting during training
+REPORT_METRICS_EVERY = 100 # for cosine similarity and other metric reporting during training
 
-REPORT_METRICS_EVERY = 250 # for cosine similarity and other metric reporting during training
 
-tracker = WandbTracker()
+def eval_model(model,device,image_reader,text_reader,start,end,batch_size,loss_type,phase="Validation"):
-
-# helpers functions
-
-def exists(val):
-    val is not None
-
-class Timer:
-    def __init__(self):
-        self.reset()
-
-    def reset(self):
-        self.last_time = time.time()
-
-    def elapsed(self):
-        return time.time() - self.last_time
-
-# functions
-
-def eval_model(model, dataloader, text_conditioned, loss_type, phase="Validation"):
     model.eval()
-
     with torch.no_grad():
         total_loss = 0.
         total_samples = 0.
 
-        for image_embeddings, text_data in tqdm(dataloader):
+        for emb_images, emb_text in zip(image_reader(batch_size=batch_size, start=start, end=end),
+                text_reader(batch_size=batch_size, start=start, end=end)):
 
-            batches = image_embeddings.shape[0]
+            emb_images_tensor = torch.tensor(emb_images[0]).to(device)
+            emb_text_tensor = torch.tensor(emb_text[0]).to(device)
 
-            input_args = dict(image_embed=image_embeddings)
+            batches = emb_images_tensor.shape[0]
-            if text_conditioned:
-                input_args = dict(**input_args, text = text_data)
-            else:
-                input_args = dict(**input_args, text_embed=text_data)
 
-            loss = model(**input_args)
+            loss = model(text_embed = emb_text_tensor, image_embed = emb_images_tensor)
 
-            total_loss += loss * batches
+            total_loss += loss.item() * batches
             total_samples += batches
 
         avg_loss = (total_loss / total_samples)
+        wandb.log({f'{phase} {loss_type}': avg_loss})
 
-        tracker.log({f'{phase} {loss_type}': avg_loss})
+def report_cosine_sims(diffusion_prior,image_reader,text_reader,train_set_size,NUM_TEST_EMBEDDINGS,device):
-
-def report_cosine_sims(diffusion_prior, dataloader, text_conditioned):
     diffusion_prior.eval()
 
     cos = nn.CosineSimilarity(dim=1, eps=1e-6)
 
-    for test_image_embeddings, text_data in tqdm(dataloader):
+    tstart = train_set_size
-
+    tend = train_set_size+NUM_TEST_EMBEDDINGS
-        # we are text conditioned, we produce an embedding from the tokenized text
-        if text_conditioned:
-            text_embedding, text_encodings, text_mask = diffusion_prior.clip.embed_text(
-                text_data)
-            text_cond = dict(text_embed=text_embedding,
-                             text_encodings=text_encodings, mask=text_mask)
-        else:
-            text_embedding = text_data
-            text_cond = dict(text_embed=text_embedding)
-
-        # make a copy of the text embeddings for shuffling
-        text_embed_shuffled = text_embedding.clone()
-
-        # roll the text to simulate "unrelated" captions
-        rolled_idx = torch.roll(torch.arange(text_embedding.shape[0]), 1)
-        text_embed_shuffled = text_embed_shuffled[rolled_idx]
-        text_embed_shuffled = text_embed_shuffled / \
-            text_embed_shuffled.norm(dim=1, keepdim=True)
-
-        if text_conditioned:
-            text_encodings_shuffled = text_encodings[rolled_idx]
-            text_mask_shuffled = text_mask[rolled_idx]
-        else:
-            text_encodings_shuffled = None
-            text_mask_shuffled = None
-
-        text_cond_shuffled = dict(text_embed=text_embed_shuffled,
-                                  text_encodings=text_encodings_shuffled, mask=text_mask_shuffled)
 
+    for embt, embi in zip(text_reader(batch_size=NUM_TEST_EMBEDDINGS, start=tstart, end=tend), 
+            image_reader(batch_size=NUM_TEST_EMBEDDINGS, start=tstart, end=tend)):
+       # make a copy of the text embeddings for shuffling
+       text_embed = torch.tensor(embt[0]).to(device)
+       text_embed_shuffled = text_embed.clone()
+        # roll the text embeddings to simulate "unrelated" captions
+       rolled_idx = torch.roll(torch.arange(NUM_TEST_EMBEDDINGS), 1)
+       text_embed_shuffled = text_embed_shuffled[rolled_idx]
+       text_embed_shuffled = text_embed_shuffled / \
+           text_embed_shuffled.norm(dim=1, keepdim=True)
+       test_text_shuffled_cond = dict(text_embed=text_embed_shuffled)
         # prepare the text embedding
-        text_embed = text_embedding / text_embedding.norm(dim=1, keepdim=True)
+       text_embed = text_embed / text_embed.norm(dim=1, keepdim=True)
-
+       test_text_cond = dict(text_embed=text_embed)
         # prepare image embeddings
-        test_image_embeddings = test_image_embeddings / \
+       test_image_embeddings = torch.tensor(embi[0]).to(device)
-            test_image_embeddings.norm(dim=1, keepdim=True)
+       test_image_embeddings = test_image_embeddings / \
-
+           test_image_embeddings.norm(dim=1, keepdim=True)
         # predict on the unshuffled text embeddings
-        predicted_image_embeddings = diffusion_prior.p_sample_loop(
+       predicted_image_embeddings = diffusion_prior.p_sample_loop(
-            test_image_embeddings.shape, text_cond)
+           (NUM_TEST_EMBEDDINGS, 768), text_cond=test_text_cond)
-        predicted_image_embeddings = predicted_image_embeddings / \
+       predicted_image_embeddings = predicted_image_embeddings / \
-            predicted_image_embeddings.norm(dim=1, keepdim=True)
+           predicted_image_embeddings.norm(dim=1, keepdim=True)
-
         # predict on the shuffled embeddings
-        predicted_unrelated_embeddings = diffusion_prior.p_sample_loop(
+       predicted_unrelated_embeddings = diffusion_prior.p_sample_loop(
-            test_image_embeddings.shape, text_cond_shuffled)
+           (NUM_TEST_EMBEDDINGS, 768), text_cond=test_text_shuffled_cond)
-        predicted_unrelated_embeddings = predicted_unrelated_embeddings / \
+       predicted_unrelated_embeddings = predicted_unrelated_embeddings / \
-            predicted_unrelated_embeddings.norm(dim=1, keepdim=True)
+           predicted_unrelated_embeddings.norm(dim=1, keepdim=True)
-
         # calculate similarities
-        original_similarity = cos(
+       original_similarity = cos(
            text_embed, test_image_embeddings).cpu().numpy()
-        predicted_similarity = cos(
+       predicted_similarity = cos(
            text_embed, predicted_image_embeddings).cpu().numpy()
-        unrelated_similarity = cos(
+       unrelated_similarity = cos(
            text_embed, predicted_unrelated_embeddings).cpu().numpy()
-        predicted_img_similarity = cos(
+       predicted_img_similarity = cos(
            test_image_embeddings, predicted_image_embeddings).cpu().numpy()
-        tracker.log({"CosineSimilarity(text_embed,image_embed)": np.mean(original_similarity),
+       wandb.log({"CosineSimilarity(text_embed,image_embed)": np.mean(original_similarity),
             "CosineSimilarity(text_embed,predicted_image_embed)":np.mean(predicted_similarity),
             "CosineSimilarity(orig_image_embed,predicted_image_embed)":np.mean(predicted_img_similarity),
             "CosineSimilarity(text_embed,predicted_unrelated_embed)": np.mean(unrelated_similarity),
             "Cosine similarity difference":np.mean(predicted_similarity - original_similarity)})
 
+def train(image_embed_dim,
+          image_embed_url,
+          text_embed_url,
+          batch_size,
+          train_percent,
+          val_percent,
+          test_percent,
+          num_epochs,
+          dp_loss_type,
+          clip,
+          dp_condition_on_text_encodings,
+          dp_timesteps,
+          dp_normformer,
+          dp_cond_drop_prob,
+          dpn_depth,
+          dpn_dim_head,
+          dpn_heads,
+          save_interval,
+          save_path,
+          device,
+          RESUME,
+          DPRIOR_PATH,
+          config,
+          wandb_entity,
+          wandb_project,
+          learning_rate=0.001,
+          max_grad_norm=0.5,
+          weight_decay=0.01,
+          dropout=0.05,
+          amp=False):
 
-@click.command()
+    # DiffusionPriorNetwork 
-@click.option("--wandb-entity", default="laion")
-@click.option("--wandb-project", default="diffusion-prior")
-@click.option("--wandb-dataset", default="LAION-5B")
-@click.option("--wandb-arch", default="DiffusionPrior")
-@click.option("--image-embed-url", default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/img_emb/")
-@click.option("--text-embed-url", default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/text_emb/")
-@click.option("--meta-url", default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/laion2B-en-metadata/")
-@click.option("--learning-rate", default=1.1e-4)
-@click.option("--weight-decay", default=6.02e-2)
-@click.option("--dropout", default=5e-2)
-@click.option("--max-grad-norm", default=0.5)
-@click.option("--num-data-points", default=250e6)
-@click.option("--batch-size", default=320)
-@click.option("--num-epochs", default=5)
-@click.option("--image-embed-dim", default=768)
-@click.option("--train-percent", default=0.9)
-@click.option("--val-percent", default=1e-7)
-@click.option("--test-percent", default=0.0999999)
-@click.option("--dpn-depth", default=12)
-@click.option("--dpn-dim-head", default=64)
-@click.option("--dpn-heads", default=12)
-@click.option("--dp-condition-on-text-encodings", default=True)
-@click.option("--dp-timesteps", default=1000)
-@click.option("--dp-normformer", default=True)
-@click.option("--dp-cond-drop-prob", default=0.1)
-@click.option("--dp-loss-type", default="l2")
-@click.option("--clip", default="ViT-L/14")
-@click.option("--amp", default=False)
-@click.option("--save-interval", default=120)
-@click.option("--save-path", default="./diffusion_prior_checkpoints")
-@click.option("--pretrained-model-path", default=None)
-@click.option("--gpu-device", default=0)
-def train(
-    wandb_entity,
-    wandb_project,
-    wandb_dataset,
-    wandb_arch,
-    image_embed_url,
-    text_embed_url,
-    meta_url,
-    learning_rate,
-    weight_decay,
-    dropout,
-    max_grad_norm,
-    num_data_points,
-    batch_size,
-    num_epochs,
-    image_embed_dim,
-    train_percent,
-    val_percent,
-    test_percent,
-    dpn_depth,
-    dpn_dim_head,
-    dpn_heads,
-    dp_condition_on_text_encodings,
-    dp_timesteps,
-    dp_normformer,
-    dp_cond_drop_prob,
-    dp_loss_type,
-    clip,
-    amp,
-    save_interval,
-    save_path,
-    pretrained_model_path,
-    gpu_device
-):
-    config = {
-        "learning_rate": learning_rate,
-        "architecture": wandb_arch,
-        "dataset": wandb_dataset,
-        "weight_decay": weight_decay,
-        "max_gradient_clipping_norm": max_grad_norm,
-        "batch_size": batch_size,
-        "epochs": num_epochs,
-        "diffusion_prior_network": {
-            "depth": dpn_depth,
-            "dim_head": dpn_dim_head,
-            "heads": dpn_heads,
-            "normformer": dp_normformer
-        },
-        "diffusion_prior": {
-            "condition_on_text_encodings": dp_condition_on_text_encodings,
-            "timesteps": dp_timesteps,
-            "cond_drop_prob": dp_cond_drop_prob,
-            "loss_type": dp_loss_type,
-            "clip": clip
-        }
-    }
-
-    # Check if DPRIOR_PATH exists(saved model path)
-
-    DPRIOR_PATH = pretrained_model_path
-    RESUME = exists(DPRIOR_PATH)
-
-    if not RESUME:
-        tracker.init(
-            entity = wandb_entity,
-            project = wandb_project,
-            config = config
-        )
-
-    # Obtain the utilized device.
-
-    has_cuda = torch.cuda.is_available()
-    if has_cuda:
-        device = torch.device(f"cuda:{gpu_device}")
-        torch.cuda.set_device(device)
-
-    # Training loop
-    # diffusion prior network
-
     prior_network = DiffusionPriorNetwork( 
-        dim = image_embed_dim,
+            dim = image_embed_dim, 
-        depth = dpn_depth,
+            depth = dpn_depth, 
-        dim_head = dpn_dim_head,
+            dim_head = dpn_dim_head, 
-        heads = dpn_heads,
+            heads = dpn_heads,
-        attn_dropout = dropout,
+            attn_dropout = dropout,
-        ff_dropout = dropout,
+            ff_dropout = dropout,
-        normformer = dp_normformer
+            normformer = dp_normformer).to(device)
-    )
     
-    # Load clip model if text-conditioning
+    # DiffusionPrior with text embeddings and image embeddings pre-computed
-    if dp_condition_on_text_encodings:
-        clip_adapter = OpenAIClipAdapter(clip)
-    else:
-        clip_adapter = None
-        
-    # diffusion prior with text embeddings and image embeddings pre-computed
-
     diffusion_prior = DiffusionPrior( 
-        net = prior_network,
+            net = prior_network, 
-        clip = clip_adapter,
+            clip = clip, 
-        image_embed_dim = image_embed_dim,
+            image_embed_dim = image_embed_dim, 
-        timesteps = dp_timesteps,
+            timesteps = dp_timesteps,
-        cond_drop_prob = dp_cond_drop_prob,
+            cond_drop_prob = dp_cond_drop_prob, 
-        loss_type = dp_loss_type,
+            loss_type = dp_loss_type, 
-        condition_on_text_encodings = dp_condition_on_text_encodings
+            condition_on_text_encodings = dp_condition_on_text_encodings).to(device)
-    )
 
     # Load pre-trained model from DPRIOR_PATH
-
     if RESUME:
-        diffusion_prior, loaded_obj = load_diffusion_model(DPRIOR_PATH, device)
+        diffusion_prior=load_diffusion_model(DPRIOR_PATH,device)   
-        tracker.init(entity = wandb_entity, project = wandb_project, config = config)
+        wandb.init( entity=wandb_entity, project=wandb_project, config=config) 
-
-    # diffusion prior trainer
-
-    trainer = DiffusionPriorTrainer(
-        diffusion_prior = diffusion_prior,
-        lr = learning_rate,
-        wd = weight_decay,
-        max_grad_norm = max_grad_norm,
-        amp = amp,
-    ).to(device)
-
-    # load optimizer and scaler
-
-    if RESUME:
-        trainer.optimizer.load_state_dict(loaded_obj['optimizer'])
-        trainer.scaler.load_state_dict(loaded_obj['scaler'])
 
     # Create save_path if it doesn't exist
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
 
-    Path(save_path).mkdir(exist_ok = True, parents = True)
+    # Get image and text embeddings from the servers
-
+    print_ribbon("Downloading embeddings - image and text")
-    # Utilize wrapper to abstract away loader logic
+    image_reader = EmbeddingReader(embeddings_folder=image_embed_url, file_format="npy")
-    print_ribbon("Downloading Embeddings")
+    text_reader  = EmbeddingReader(embeddings_folder=text_embed_url, file_format="npy")
-    loader_args = dict(text_conditioned=dp_condition_on_text_encodings, batch_size=batch_size, num_data_points=num_data_points,
+    num_data_points = text_reader.count
-                       train_split=train_percent, eval_split=val_percent, device=device, img_url=image_embed_url)
-
-    if dp_condition_on_text_encodings:
-        loader_args = dict(**loader_args, meta_url=meta_url)
-    else:
-        loader_args = dict(**loader_args, txt_url=text_embed_url)
-
-    train_loader, eval_loader, test_loader = make_splits(**loader_args)
 
     ### Training code ###
-
+    scaler = GradScaler(enabled=amp)
-    step = 1 
+    optimizer = get_optimizer(diffusion_prior.net.parameters(), wd=weight_decay, lr=learning_rate)
-    timer = Timer()
     epochs = num_epochs
 
+    step = 0
+    t = time.time()
+
+    train_set_size = int(train_percent*num_data_points)
+    val_set_size = int(val_percent*num_data_points)
+    eval_start = train_set_size
+
     for _ in range(epochs):
 
-        for image, text in tqdm(train_loader):
+        for emb_images,emb_text in zip(image_reader(batch_size=batch_size, start=0, end=train_set_size),
-            
+                text_reader(batch_size=batch_size, start=0, end=train_set_size)):
+
             diffusion_prior.train()
             
-            input_args = dict(image_embed=image)
+            emb_images_tensor = torch.tensor(emb_images[0]).to(device)
-            if dp_condition_on_text_encodings:
+            emb_text_tensor = torch.tensor(emb_text[0]).to(device)
-                input_args = dict(**input_args, text = text)
-            else:
-                input_args = dict(**input_args, text_embed=text)
 
-            loss = trainer(**input_args)
+            with autocast(enabled=amp):
+                loss = diffusion_prior(text_embed = emb_text_tensor,image_embed = emb_images_tensor)
+                scaler.scale(loss).backward()
 
             # Samples per second
-
+            step+=1
-            samples_per_sec = batch_size * step / timer.elapsed()
+            samples_per_sec = batch_size*step/(time.time()-t)
-
             # Save checkpoint every save_interval minutes
-            if(int(timer.elapsed()) >= 60 * save_interval):
+            if(int(time.time()-t) >= 60*save_interval):
-                timer.reset()
+                t = time.time()
 
                 save_diffusion_model(
                     save_path,
                     diffusion_prior,
-                    trainer.optimizer,
+                    optimizer,
-                    trainer.scaler,
+                    scaler,
                     config,
                     image_embed_dim)
 
             # Log to wandb
-            tracker.log({"Training loss": loss,
+            wandb.log({"Training loss": loss.item(),
                         "Steps": step,
                         "Samples per second": samples_per_sec})
             # Log cosineSim(text_embed,predicted_image_embed) - cosineSim(text_embed,image_embed)
             # Use NUM_TEST_EMBEDDINGS samples from the test set each time
             # Get embeddings from the most recently saved model
             if(step % REPORT_METRICS_EVERY) == 0:
-                report_cosine_sims(diffusion_prior, eval_loader, dp_condition_on_text_encodings)
+                report_cosine_sims(diffusion_prior,
+                        image_reader,
+                        text_reader,
+                        train_set_size,
+                        NUM_TEST_EMBEDDINGS,
+                        device)
                 ### Evaluate model(validation run) ###
-                eval_model(diffusion_prior, eval_loader, dp_condition_on_text_encodings, dp_loss_type, phase="Validation")
+                eval_model(diffusion_prior,
+                        device,
+                        image_reader,
+                        text_reader,
+                        eval_start,
+                        eval_start+NUM_TEST_EMBEDDINGS,
+                        NUM_TEST_EMBEDDINGS,
+                        dp_loss_type,
+                        phase="Validation")
 
-            step += 1
+            scaler.unscale_(optimizer)
-            trainer.update()
+            nn.utils.clip_grad_norm_(diffusion_prior.parameters(), max_grad_norm)
+
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad()
 
     ### Test run ###
-    eval_model(diffusion_prior, test_loader, dp_condition_on_text_encodings, dp_loss_type, phase="Test")
+    test_set_size = int(test_percent*train_set_size) 
+    start=train_set_size+val_set_size
+    end=num_data_points
+    eval_model(diffusion_prior,device,image_reader,text_reader,start,end,batch_size,dp_loss_type,phase="Test")
 
+def main():
+    parser = argparse.ArgumentParser()
+    # Logging
+    parser.add_argument("--wandb-entity", type=str, default="laion")
+    parser.add_argument("--wandb-project", type=str, default="diffusion-prior")
+    parser.add_argument("--wandb-dataset", type=str, default="LAION-5B")
+    parser.add_argument("--wandb-arch", type=str, default="DiffusionPrior")
+    # URLs for embeddings 
+    parser.add_argument("--image-embed-url", type=str, default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/img_emb/")
+    parser.add_argument("--text-embed-url", type=str, default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/text_emb/")
+    # Hyperparameters
+    parser.add_argument("--learning-rate", type=float, default=1.1e-4)
+    parser.add_argument("--weight-decay", type=float, default=6.02e-2)
+    parser.add_argument("--dropout", type=float, default=5e-2)
+    parser.add_argument("--max-grad-norm", type=float, default=0.5)
+    parser.add_argument("--batch-size", type=int, default=10**4)
+    parser.add_argument("--num-epochs", type=int, default=5)
+    # Image embed dimension
+    parser.add_argument("--image-embed-dim", type=int, default=768)
+    # Train-test split
+    parser.add_argument("--train-percent", type=float, default=0.7)
+    parser.add_argument("--val-percent", type=float, default=0.2)
+    parser.add_argument("--test-percent", type=float, default=0.1)
+    # LAION training(pre-computed embeddings)
+    # DiffusionPriorNetwork(dpn) parameters
+    parser.add_argument("--dpn-depth", type=int, default=6)
+    parser.add_argument("--dpn-dim-head", type=int, default=64)
+    parser.add_argument("--dpn-heads", type=int, default=8)
+    # DiffusionPrior(dp) parameters
+    parser.add_argument("--dp-condition-on-text-encodings", type=bool, default=False)
+    parser.add_argument("--dp-timesteps", type=int, default=100)
+    parser.add_argument("--dp-normformer", type=bool, default=False)
+    parser.add_argument("--dp-cond-drop-prob", type=float, default=0.1)
+    parser.add_argument("--dp-loss-type", type=str, default="l2")
+    parser.add_argument("--clip", type=str, default=None)
+    parser.add_argument("--amp", type=bool, default=False)
+    # Model checkpointing interval(minutes)
+    parser.add_argument("--save-interval", type=int, default=30)
+    parser.add_argument("--save-path", type=str, default="./diffusion_prior_checkpoints")
+    # Saved model path 
+    parser.add_argument("--pretrained-model-path", type=str, default=None)
+
+    args = parser.parse_args()
+
+    config = ({"learning_rate": args.learning_rate,
+        "architecture": args.wandb_arch,
+        "dataset": args.wandb_dataset,
+        "weight_decay":args.weight_decay,
+        "max_gradient_clipping_norm":args.max_grad_norm,
+        "batch_size":args.batch_size,
+        "epochs": args.num_epochs,
+        "diffusion_prior_network":{"depth":args.dpn_depth,
+        "dim_head":args.dpn_dim_head,
+        "heads":args.dpn_heads,
+        "normformer":args.dp_normformer},
+        "diffusion_prior":{"condition_on_text_encodings": args.dp_condition_on_text_encodings,
+        "timesteps": args.dp_timesteps,
+        "cond_drop_prob":args.dp_cond_drop_prob,
+        "loss_type":args.dp_loss_type,
+        "clip":args.clip}
+        })
+
+    RESUME = False
+    # Check if DPRIOR_PATH exists(saved model path)
+    DPRIOR_PATH = args.pretrained_model_path
+    if(DPRIOR_PATH is not None):
+        RESUME = True
+    else:
+        wandb.init(
+          entity=args.wandb_entity,
+          project=args.wandb_project,
+          config=config)
+
+    # Obtain the utilized device.
+
+    has_cuda = torch.cuda.is_available()
+    if has_cuda:
+        device = torch.device("cuda:0")
+        torch.cuda.set_device(device)
+
+    # Training loop
+    train(args.image_embed_dim,
+          args.image_embed_url,
+          args.text_embed_url,
+          args.batch_size,
+          args.train_percent,
+          args.val_percent,
+          args.test_percent,
+          args.num_epochs,
+          args.dp_loss_type,
+          args.clip,
+          args.dp_condition_on_text_encodings,
+          args.dp_timesteps,
+          args.dp_normformer,
+          args.dp_cond_drop_prob,
+          args.dpn_depth,
+          args.dpn_dim_head,
+          args.dpn_heads,
+          args.save_interval,
+          args.save_path,
+          device,
+          RESUME,
+          DPRIOR_PATH,
+          config,
+          args.wandb_entity,
+          args.wandb_project,
+          args.learning_rate,
+          args.max_grad_norm,
+          args.weight_decay,
+          args.dropout,
+          args.amp)
 
 if __name__ == "__main__":
-    train()
+  main()