trainer classes now takes care of auto-casting numpy to torch tensors, and setting correct device based on model parameter devices

README.md

@@ -814,8 +814,8 @@ clip = CLIP(
 
 # mock data
 
-text = torch.randint(0, 49408, (4, 256)).cuda()
+text = torch.randint(0, 49408, (32, 256)).cuda()
-images = torch.randn(4, 3, 256, 256).cuda()
+images = torch.randn(32, 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)
+loss = diffusion_prior_trainer(text, images, max_batch_size = 4)
 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
@@ -895,14 +895,14 @@ dataset = ImageEmbeddingDataset(
 )
 ```
 
-## Scripts
+### Scripts (wip)
 
-### Using the `train_diffusion_prior.py` script
+#### `train_diffusion_prior.py`
 
 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
@@ -910,58 +910,49 @@ $ 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
 
-### Sample wandb run log
+#### Loading and Saving the DiffusionPrior model
-
-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. 
 
-## from dalle2_pytorch.train import load_diffusion_model, save_diffusion_model
+```python
+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)
@@ -1021,6 +1012,8 @@ Once built, images will be saved to the same directory the command is invoked
 - [ ] 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)
 
 ## Citations
 

dalle2_pytorch/dalle2_pytorch.py

@@ -1988,8 +1988,7 @@ class Decoder(BaseGaussianDiffusion):
                 image_size = vae.get_encoded_fmap_size(image_size)
                 shape = (batch_size, vae.encoded_dim, image_size, image_size)
 
-                if exists(lowres_cond_img):
+                lowres_cond_img = maybe(vae.encode)(lowres_cond_img)
-                    lowres_cond_img = vae.encode(lowres_cond_img)
 
                 img = self.p_sample_loop(
                     unet,
@@ -2063,9 +2062,7 @@ 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)
 

dalle2_pytorch/train.py

@@ -1,7 +1,7 @@
 import time
 import copy
 from math import ceil
-from functools import partial
+from functools import partial, wraps
 from collections.abc import Iterable
 
 import torch
@@ -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
 
+import numpy as np
+
 # helper functions
 
 def exists(val):
@@ -45,6 +47,29 @@ 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):
@@ -66,18 +91,27 @@ def split(t, split_size = None):
 
     return TypeError
 
-def split_args_and_kwargs(x, *args, split_size = None, **kwargs):
+def find_first(cond, arr):
-    batch_size = len(x)
+    for el in arr:
+        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)
-    chunk_size = ceil(batch_size / split_size)
+    num_chunks = 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,) * chunk_size) 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,) * num_chunks) 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)):
@@ -117,7 +151,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
+    return diffusion_prior, loaded_obj
 
 def save_diffusion_model(save_path, model, optimizer, scaler, config, image_embed_dim):
     # Saving State Dict
@@ -245,10 +279,12 @@ class DiffusionPriorTrainer(nn.Module):
         self.step += 1
 
     @torch.inference_mode()
+    @cast_torch_tensor
     def p_sample_loop(self, *args, **kwargs):
         return self.ema_diffusion_prior.ema_model.p_sample_loop(*args, **kwargs)
 
     @torch.inference_mode()
+    @cast_torch_tensor
     def sample(self, *args, **kwargs):
         return self.ema_diffusion_prior.ema_model.sample(*args, **kwargs)
 
@@ -256,22 +292,24 @@ class DiffusionPriorTrainer(nn.Module):
     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(x, *args, split_size = max_batch_size, **kwargs):
+        for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*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
 
             total_loss += loss.item()
-            self.scaler.scale(loss).backward()
+
+            if self.training:
+                self.scaler.scale(loss).backward()
 
         return total_loss
 
@@ -367,6 +405,7 @@ class DecoderTrainer(nn.Module):
         self.step += 1
 
     @torch.no_grad()
+    @cast_torch_tensor
     def sample(self, *args, **kwargs):
         if self.use_ema:
             trainable_unets = self.decoder.unets
@@ -383,22 +422,24 @@ class DecoderTrainer(nn.Module):
 
         return output
 
+    @cast_torch_tensor
     def forward(
         self,
-        x,
+        *args,
-        *,
         unet_number,
         max_batch_size = None,
         **kwargs
     ):
         total_loss = 0.
 
-        for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(x, split_size = max_batch_size, **kwargs):
+        for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, 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
 
             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.30',
+  version = '0.2.34',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',
@@ -30,6 +30,7 @@ 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,27 +1,43 @@
-import os
+from pathlib import Path
+import click
 import math
 import time
-import argparse
 import numpy as np
 
 import torch
 from torch import nn
-from torch.cuda.amp import autocast, GradScaler
 
 from dalle2_pytorch import DiffusionPrior, DiffusionPriorNetwork
-from dalle2_pytorch.train import load_diffusion_model, save_diffusion_model, print_ribbon
+from dalle2_pytorch.train import DiffusionPriorTrainer, load_diffusion_model, save_diffusion_model, print_ribbon
-from dalle2_pytorch.optimizer import get_optimizer
 from dalle2_pytorch.trackers import ConsoleTracker, WandbTracker
 
 from embedding_reader import EmbeddingReader
 
 from tqdm import tqdm
 
+# constants
+
 NUM_TEST_EMBEDDINGS = 100 # for cosine similarity reporting during training
 REPORT_METRICS_EVERY = 100 # for cosine similarity and other metric reporting during training
 
 tracker = WandbTracker()
 
+# 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,device,image_reader,text_reader,start,end,batch_size,loss_type,phase="Validation"):
     model.eval()
     with torch.no_grad():
@@ -95,79 +111,171 @@ def report_cosine_sims(diffusion_prior,image_reader,text_reader,train_set_size,N
             "CosineSimilarity(text_embed,predicted_unrelated_embed)": np.mean(unrelated_similarity),
             "Cosine similarity difference":np.mean(predicted_similarity - original_similarity)})
 
-def train(image_embed_dim,
+@click.command()
-          image_embed_url,
+@click.option("--wandb-entity", default="laion")
-          text_embed_url,
+@click.option("--wandb-project", default="diffusion-prior")
-          batch_size,
+@click.option("--wandb-dataset", default="LAION-5B")
-          train_percent,
+@click.option("--wandb-arch", default="DiffusionPrior")
-          val_percent,
+@click.option("--image-embed-url", default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/img_emb/")
-          test_percent,
+@click.option("--text-embed-url", default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/text_emb/")
-          num_epochs,
+@click.option("--learning-rate", default=1.1e-4)
-          dp_loss_type,
+@click.option("--weight-decay", default=6.02e-2)
-          clip,
+@click.option("--dropout", default=5e-2)
-          dp_condition_on_text_encodings,
+@click.option("--max-grad-norm", default=0.5)
-          dp_timesteps,
+@click.option("--batch-size", default=10**4)
-          dp_normformer,
+@click.option("--num-epochs", default=5)
-          dp_cond_drop_prob,
+@click.option("--image-embed-dim", default=768)
-          dpn_depth,
+@click.option("--train-percent", default=0.7)
-          dpn_dim_head,
+@click.option("--val-percent", default=0.2)
-          dpn_heads,
+@click.option("--test-percent", default=0.1)
-          save_interval,
+@click.option("--dpn-depth", default=6)
-          save_path,
+@click.option("--dpn-dim-head", default=64)
-          device,
+@click.option("--dpn-heads", default=8)
-          RESUME,
+@click.option("--dp-condition-on-text-encodings", default=False)
-          DPRIOR_PATH,
+@click.option("--dp-timesteps", default=100)
-          config,
+@click.option("--dp-normformer", default=False)
-          wandb_entity,
+@click.option("--dp-cond-drop-prob", default=0.1)
-          wandb_project,
+@click.option("--dp-loss-type", default="l2")
-          learning_rate=0.001,
+@click.option("--clip", default=None)
-          max_grad_norm=0.5,
+@click.option("--amp", default=False)
-          weight_decay=0.01,
+@click.option("--save-interval", default=30)
-          dropout=0.05,
+@click.option("--save-path", default="./diffusion_prior_checkpoints")
-          amp=False):
+@click.option("--pretrained-model-path", default=None)
+def train(
+    wandb_entity,
+    wandb_project,
+    wandb_dataset,
+    wandb_arch,
+    image_embed_url,
+    text_embed_url,
+    learning_rate,
+    weight_decay,
+    dropout,
+    max_grad_norm,
+    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
+):
+    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 = args.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("cuda:0")
+        torch.cuda.set_device(device)
+
+    # Training loop
+    # diffusion prior network
 
-    # DiffusionPriorNetwork 
     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).to(device)
+        normformer = dp_normformer
+    )
     
-    # DiffusionPrior with text embeddings and image embeddings pre-computed
+    # diffusion prior with text embeddings and image embeddings pre-computed
+
     diffusion_prior = DiffusionPrior( 
-            net = prior_network, 
+        net = prior_network,
-            clip = clip, 
+        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).to(device)
+        condition_on_text_encodings = dp_condition_on_text_encodings
+    )
 
     # Load pre-trained model from DPRIOR_PATH
+
     if RESUME:
-        diffusion_prior=load_diffusion_model(DPRIOR_PATH,device)   
+        diffusion_prior, loaded_obj = load_diffusion_model(DPRIOR_PATH, device)
-        wandb.init( entity=wandb_entity, project=wandb_project, config=config) 
+        tracker.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")
     image_reader = EmbeddingReader(embeddings_folder=image_embed_url, file_format="npy")
     text_reader  = EmbeddingReader(embeddings_folder=text_embed_url, file_format="npy")
     num_data_points = text_reader.count
 
     ### Training code ###
-    scaler = GradScaler(enabled=amp)
-    optimizer = get_optimizer(diffusion_prior.net.parameters(), wd=weight_decay, lr=learning_rate)
-    epochs = num_epochs
 
-    step = 0
+    timer = Timer()
-    t = time.time()
+    epochs = num_epochs
 
     train_set_size = int(train_percent*num_data_points)
     val_set_size = int(val_percent*num_data_points)
@@ -178,27 +286,26 @@ def train(image_embed_dim,
         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()
+            trainer.train()
             
             emb_images_tensor = torch.tensor(emb_images[0]).to(device)
             emb_text_tensor = torch.tensor(emb_text[0]).to(device)
 
-            with autocast(enabled=amp):
+            loss = trainer(text_embed = emb_text_tensor, image_embed = emb_images_tensor)
-                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/(time.time()-t)
+            samples_per_sec = batch_size * step / timer.elapsed()
+
             # Save checkpoint every save_interval minutes
-            if(int(time.time()-t) >= 60*save_interval):
+            if(int(timer.elapsed()) >= 60 * save_interval):
-                t = time.time()
+                timer.reset()
 
                 save_diffusion_model(
                     save_path,
                     diffusion_prior,
-                    optimizer,
+                    trainer.optimizer,
-                    scaler,
+                    trainer.scaler,
                     config,
                     image_embed_dim)
 
@@ -227,130 +334,13 @@ def train(image_embed_dim,
                         dp_loss_type,
                         phase="Validation")
 
-            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 ###
     test_set_size = int(test_percent*train_set_size) 
-    start=train_set_size+val_set_size
+    start = train_set_size+val_set_size
-    end=num_data_points
+    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:
-        tracker.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__":
-    main()
+    train()