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

README.md

@@ -814,8 +814,8 @@ clip = CLIP(
 
 # mock data
 
-text = torch.randint(0, 49408, (32, 256)).cuda()
-images = torch.randn(32, 3, 256, 256).cuda()
+text = torch.randint(0, 49408, (4, 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
@@ -1007,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
@@ -1015,6 +1014,7 @@ 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

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

@@ -66,24 +66,15 @@ def split(t, split_size = None):
 
     return TypeError
 
-def find_first(cond, arr):
-    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)
+def split_args_and_kwargs(x, *args, split_size = None, **kwargs):
+    batch_size = len(x)
     split_size = default(split_size, batch_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,) * chunk_size) for arg in all_args]
@@ -126,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
@@ -237,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)
@@ -251,8 +240,6 @@ class DiffusionPriorTrainer(nn.Module):
         if self.use_ema:
             self.ema_diffusion_prior.update()
 
-        self.step += 1
-
     @torch.inference_mode()
     def p_sample_loop(self, *args, **kwargs):
         return self.ema_diffusion_prior.ema_model.p_sample_loop(*args, **kwargs)
@@ -267,17 +254,18 @@ class DiffusionPriorTrainer(nn.Module):
 
     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()
 
@@ -340,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])
@@ -372,8 +358,6 @@ class DecoderTrainer(nn.Module):
             ema_unet = self.ema_unets[index]
             ema_unet.update()
 
-        self.step += 1
-
     @torch.no_grad()
     def sample(self, *args, **kwargs):
         if self.use_ema:
@@ -393,18 +377,19 @@ class DecoderTrainer(nn.Module):
 
     def forward(
         self,
-        *args,
+        x,
+        *,
         unet_number,
         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, 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()
 

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.2.31',
+  version = '0.2.28',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',

train_diffusion_prior.py

@@ -1,42 +1,24 @@
-from pathlib import Path
-import click
+import os
 import math
-import time
+import argparse
 import numpy as np
 
 import torch
 from torch import nn
-
-from dalle2_pytorch import DiffusionPrior, DiffusionPriorNetwork
-from dalle2_pytorch.train 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
 
-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()
@@ -58,7 +40,7 @@ def eval_model(model,device,image_reader,text_reader,start,end,batch_size,loss_t
             total_samples += batches
 
         avg_loss = (total_loss / total_samples)
-        tracker.log({f'{phase} {loss_type}': avg_loss})
+        wandb.log({f'{phase} {loss_type}': avg_loss})
 
 def report_cosine_sims(diffusion_prior,image_reader,text_reader,train_set_size,NUM_TEST_EMBEDDINGS,device):
     diffusion_prior.eval()
@@ -105,7 +87,7 @@ def report_cosine_sims(diffusion_prior,image_reader,text_reader,train_set_size,N
            text_embed, predicted_unrelated_embeddings).cpu().numpy()
        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),
@@ -142,68 +124,49 @@ def train(image_embed_dim,
           dropout=0.05,
           amp=False):
 
-    # diffusion prior network
-
+    # DiffusionPriorNetwork 
     prior_network = DiffusionPriorNetwork( 
-        dim = image_embed_dim,
-        depth = dpn_depth,
-        dim_head = dpn_dim_head,
-        heads = dpn_heads,
-        attn_dropout = dropout,
-        ff_dropout = dropout,
-        normformer = dp_normformer
-    )
+            dim = image_embed_dim, 
+            depth = dpn_depth, 
+            dim_head = dpn_dim_head, 
+            heads = dpn_heads,
+            attn_dropout = dropout,
+            ff_dropout = dropout,
+            normformer = dp_normformer).to(device)
     
-    # diffusion prior with text embeddings and image embeddings pre-computed
-
+    # DiffusionPrior with text embeddings and image embeddings pre-computed
     diffusion_prior = DiffusionPrior( 
-        net = prior_network,
-        clip = clip,
-        image_embed_dim = image_embed_dim,
-        timesteps = dp_timesteps,
-        cond_drop_prob = dp_cond_drop_prob,
-        loss_type = dp_loss_type,
-        condition_on_text_encodings = dp_condition_on_text_encodings
-    )
+            net = prior_network, 
+            clip = clip, 
+            image_embed_dim = image_embed_dim, 
+            timesteps = dp_timesteps,
+            cond_drop_prob = dp_cond_drop_prob, 
+            loss_type = dp_loss_type, 
+            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)
-        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'])
+        diffusion_prior=load_diffusion_model(DPRIOR_PATH,device)   
+        wandb.init( entity=wandb_entity, project=wandb_project, config=config) 
 
     # Create save_path if it doesn't exist
-
-    Path(save_path).mkdir(exist_ok = True, parents = True)
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
 
     # 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 ###
-
-    timer = Timer()
+    scaler = GradScaler(enabled=amp)
+    optimizer = get_optimizer(diffusion_prior.net.parameters(), wd=weight_decay, lr=learning_rate)
     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
@@ -213,31 +176,32 @@ 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)):
 
-            trainer.train()
+            diffusion_prior.train()
             
             emb_images_tensor = torch.tensor(emb_images[0]).to(device)
             emb_text_tensor = torch.tensor(emb_text[0]).to(device)
 
-            loss = trainer(text_embed = emb_text_tensor, image_embed = emb_images_tensor)
+            with autocast(enabled=amp):
+                loss = diffusion_prior(text_embed = emb_text_tensor,image_embed = emb_images_tensor)
+                scaler.scale(loss).backward()
 
             # Samples per second
-
-            samples_per_sec = batch_size * step / timer.elapsed()
-
+            step+=1
+            samples_per_sec = batch_size*step/(time.time()-t)
             # Save checkpoint every save_interval minutes
-            if(int(timer.elapsed()) >= 60 * save_interval):
-                timer.reset()
+            if(int(time.time()-t) >= 60*save_interval):
+                t = time.time()
 
                 save_diffusion_model(
                     save_path,
                     diffusion_prior,
-                    trainer.optimizer,
-                    trainer.scaler,
+                    optimizer,
+                    scaler,
                     config,
                     image_embed_dim)
 
             # Log to wandb
-            tracker.log({"Training loss": loss.item(),
+            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)
@@ -261,109 +225,91 @@ def train(image_embed_dim,
                         dp_loss_type,
                         phase="Validation")
 
-            trainer.update()
+            scaler.unscale_(optimizer)
+            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
-    end = num_data_points
+    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")
 
-@click.command()
-@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("--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("--batch-size", default=10**4)
-@click.option("--num-epochs", default=5)
-@click.option("--image-embed-dim", default=768)
-@click.option("--train-percent", default=0.7)
-@click.option("--val-percent", default=0.2)
-@click.option("--test-percent", default=0.1)
-@click.option("--dpn-depth", default=6)
-@click.option("--dpn-dim-head", default=64)
-@click.option("--dpn-heads", default=8)
-@click.option("--dp-condition-on-text-encodings", default=False)
-@click.option("--dp-timesteps", default=100)
-@click.option("--dp-normformer", default=False)
-@click.option("--dp-cond-drop-prob", default=0.1)
-@click.option("--dp-loss-type", default="l2")
-@click.option("--clip", default=None)
-@click.option("--amp", default=False)
-@click.option("--save-interval", default=30)
-@click.option("--save-path", default="./diffusion_prior_checkpoints")
-@click.option("--pretrained-model-path", default=None)
-def main(
-    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
-        }
-    }
+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
-    RESUME = exists(DPRIOR_PATH)
-
-    if not RESUME:
-        tracker.init(
-            entity = wandb_entity,
-            project = wandb_project,
-            config = config
-        )
+    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.
 
@@ -373,36 +319,36 @@ def main(
         torch.cuda.set_device(device)
 
     # Training loop
-    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,
+    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,
-          wandb_entity,
-          wandb_project,
-          learning_rate,
-          max_grad_norm,
-          weight_decay,
-          dropout,
-          amp)
+          args.wandb_entity,
+          args.wandb_project,
+          args.learning_rate,
+          args.max_grad_norm,
+          args.weight_decay,
+          args.dropout,
+          args.amp)
 
 if __name__ == "__main__":
-    main()
+  main()