fix everything and make sure it runs end to end, document everything in readme for public

dalle2_pytorch/__init__.py

@@ -1,6 +1,2 @@
 from dalle2_pytorch.dalle2_pytorch import DALLE2, DiffusionPriorNetwork, DiffusionPrior, Unet, Decoder
-from dalle2_pytorch.dalle2_pytorch import OpenAIClipAdapter
-from dalle2_pytorch.train import DecoderTrainer, DiffusionPriorTrainer
-
-from dalle2_pytorch.vqgan_vae import VQGanVAE
 from x_clip import CLIP

dalle2_pytorch/cli.py

@@ -1,52 +1,9 @@
 import click
-import torch
-import torchvision.transforms as T
-from functools import reduce
-from pathlib import Path
-
-from dalle2_pytorch import DALLE2, Decoder, DiffusionPrior
-
-def safeget(dictionary, keys, default = None):
-    return reduce(lambda d, key: d.get(key, default) if isinstance(d, dict) else default, keys.split('.'), dictionary)
-
-def simple_slugify(text, max_length = 255):
-    return text.replace("-", "_").replace(",", "").replace(" ", "_").replace("|", "--").strip('-_')[:max_length]
-
-def get_pkg_version():
-    from pkg_resources import get_distribution
-    return get_distribution('dalle2_pytorch').version
 
 def main():
     pass
 
 @click.command()
-@click.option('--model', default = './dalle2.pt', help = 'path to trained DALL-E2 model')
-@click.option('--cond_scale', default = 2, help = 'conditioning scale (classifier free guidance) in decoder')
 @click.argument('text')
-def dream(
-    model,
-    cond_scale,
-    text
-):
-    model_path = Path(model)
-    full_model_path = str(model_path.resolve())
-    assert model_path.exists(), f'model not found at {full_model_path}'
-    loaded = torch.load(str(model_path))
-
-    version = safeget(loaded, 'version')
-    print(f'loading DALL-E2 from {full_model_path}, saved at version {version} - current package version is {get_pkg_version()}')
-
-    prior_init_params = safeget(loaded, 'init_params.prior')
-    decoder_init_params = safeget(loaded, 'init_params.decoder')
-    model_params = safeget(loaded, 'model_params')
-
-    prior = DiffusionPrior(**prior_init_params)
-    decoder = Decoder(**decoder_init_params)
-
-    dalle2 = DALLE2(prior, decoder)
-    dalle2.load_state_dict(model_params)
-
-    image = dalle2(text, cond_scale = cond_scale)
-
-    pil_image = T.ToPILImage()(image)
-    return pil_image.save(f'./{simple_slugify(text)}.png')
+def dream(text):
+    return image

dalle2_pytorch/dataloaders/__init__.py

@@ -1 +0,0 @@
-from dalle2_pytorch.dataloaders.decoder_loader import ImageEmbeddingDataset, create_image_embedding_dataloader

dalle2_pytorch/dataloaders/decoder_loader.py

@@ -1,170 +0,0 @@
-import os
-import webdataset as wds
-import torch
-import numpy as np
-import fsspec
-
-def get_shard(filename):
-    """
-    Filenames with shards in them have a consistent structure that we can take advantage of
-    Standard structure: path/to/file/prefix_string_00001.ext
-    """
-    try:
-        return filename.split("_")[-1].split(".")[0]
-    except ValueError:
-        raise RuntimeError(f"Could not find shard for filename {filename}")
-
-def get_example_file(fs, path, file_format):
-    """
-    Given a file system and a file extension, return the example file
-    """
-    return fs.glob(os.path.join(path, f"*.{file_format}"))[0]
-
-def embedding_inserter(samples, embeddings_url, shard_width, handler=wds.handlers.reraise_exception):
-    """Given a datum of {"__key__": str, "__url__": str, ...} adds the cooresponding embedding and yields"""
-    previous_tar_url = None
-    current_embeddings = None
-    # Get a reference to an abstract file system where the embeddings are stored
-    embeddings_fs, embeddings_path = fsspec.core.url_to_fs(embeddings_url)
-    example_embedding_file = get_example_file(embeddings_fs, embeddings_path, "npy")
-    example_embedding_shard = get_shard(example_embedding_file)
-    emb_shard_width = len(example_embedding_shard)
-    # Easier to get the basename without the shard once than search through for the correct file every time
-    embedding_file_basename = '_'.join(example_embedding_file.split("_")[:-1]) + "_"
-
-    def load_corresponding_embeds(tar_url):
-      """Finds and reads the npy files that contains embeddings for the given webdataset tar"""
-      shard = int(tar_url.split("/")[-1].split(".")[0])
-      embedding_url = embedding_file_basename + str(shard).zfill(emb_shard_width) + '.npy'
-      with embeddings_fs.open(embedding_url) as f:
-        data = np.load(f)
-      return torch.from_numpy(data)
-
-    for sample in samples:
-        try:
-            tar_url = sample["__url__"]
-            key = sample["__key__"]
-            if tar_url != previous_tar_url:
-                # If the tar changed, we need to download new embeddings
-                # This means if we shuffle before inserting it will load many more files than we expect and be very inefficient.
-                previous_tar_url = tar_url
-                current_embeddings = load_corresponding_embeds(tar_url)
-                
-            embedding_index = int(key[shard_width:])
-            sample["npy"] = current_embeddings[embedding_index]
-            yield sample
-        except Exception as exn:  # From wds implementation
-            if handler(exn):
-                continue
-            else:
-                break
-insert_embedding = wds.filters.pipelinefilter(embedding_inserter)
-
-def verify_keys(samples, handler=wds.handlers.reraise_exception):
-    """
-    Requires that both the image and embedding are present in the sample
-    This is important to do as a user may forget they do not have embeddings in their webdataset and neglect to add them using the embedding_folder_url parameter.
-    """
-    for sample in samples:
-        try:
-            assert "jpg" in sample, f"Sample {sample['__key__']} missing image"
-            assert "npy" in sample, f"Sample {sample['__key__']} missing embedding. Did you set embedding_folder_url?"
-            yield sample
-        except Exception as exn:  # From wds implementation
-            if handler(exn):
-                continue
-            else:
-                break
-
-class ImageEmbeddingDataset(wds.DataPipeline, wds.compat.FluidInterface):
-    """
-    A fluid interface wrapper for DataPipline that returns image embedding pairs
-    Reads embeddings as npy files from the webdataset if they exist. If embedding_folder_url is set, they will be inserted in from the alternate source.
-    """
-
-    def __init__(
-            self,
-            urls,
-            embedding_folder_url=None,
-            shard_width=None,
-            handler=wds.handlers.reraise_exception,
-            resample=False,
-            shuffle_shards=True
-    ):
-        """
-        Modeled directly off of the WebDataset constructor
-
-        :param urls: A url pointing to the tar files of the webdataset formatted as /path/to/webdataset/{0000..9999}.tar
-        :param embedding_folder_url: Required if webdataset does not contain embeddings. A url pointing to the npy files of the embeddings. Should have the same number of shards as the webdataset.
-            Webdataset image keys should align with the index of the embedding. This means missing image indices must have a corresponding embedding of all zeros.
-        :param shard_width: The number of digits in the shard number. This is used to align the embedding index with the image index.
-            For example, if a file in the webdataset shard 3 is named 0003039.jpg, we know the shard with this 4 and the last three digits are the index.
-        :param handler: A webdataset handler.
-        :param resample: If true, resample webdataset shards with replacement. You need to set your own epoch size if this is true since it will resample infinitely.
-        :param shuffle_shards: If true, shuffle the shards before resampling. This cannot be true if resample is true.
-        """
-        super().__init__()
-        # Add the shardList and randomize or resample if requested
-        if resample:
-            assert not shuffle_shards, "Cannot both resample and shuffle"
-            self.append(wds.ResampledShards(urls))
-        else:
-            self.append(wds.SimpleShardList(urls))
-            if shuffle_shards:
-                self.append(wds.filters.shuffle(1000))
-
-        self.append(wds.split_by_node)
-        self.append(wds.split_by_worker)
-
-        self.append(wds.tarfile_to_samples(handler=handler))
-        self.append(wds.decode("torchrgb"))
-        if embedding_folder_url is not None:
-            assert shard_width is not None, "Reading embeddings separately requires shard length to be given"
-            self.append(insert_embedding(embeddings_url=embedding_folder_url, shard_width=shard_width, handler=handler))
-        self.append(verify_keys)
-        self.append(wds.to_tuple("jpg", "npy"))
-
-def create_image_embedding_dataloader(
-    tar_url,
-    num_workers,
-    batch_size,
-    embeddings_url=None,
-    shard_width=None,
-    shuffle_num = None,
-    shuffle_shards = True,
-    resample_shards = False, 
-    handler=wds.handlers.warn_and_continue
-):
-    """
-    Convenience function to create an image embedding dataseta and dataloader in one line
-
-    :param tar_url: A url pointing to the tar files of the webdataset formatted as /path/to/webdataset/{0000..9999}.tar
-    :param num_workers: The number of workers to use for the dataloader
-    :param batch_size: The batch size to use for the dataloader
-    :param embeddings_url: Required if webdataset does not contain embeddings. A url pointing to the npy files of the embeddings. Should have the same number of shards as the webdataset.
-        Webdataset image keys should align with the index of the embedding. This means missing image indices must have a corresponding embedding of all zeros.
-    :param shard_width: The number of digits in the shard number. This is used to align the embedding index with the image index.
-        For example, if a file in the webdataset shard 3 is named 0003039.jpg, we know the shard width is 4 and the last three digits are the index.
-    :param shuffle_num: If not None, shuffle the dataset with this size buffer after sampling.
-    :param shuffle_shards: If true, shuffle the shards before sampling. This cannot be true if resample is true.
-    :param resample_shards: If true, resample webdataset shards with replacement. You need to set your own epoch size if this is true since it will resample infinitely.
-    :param handler: A webdataset handler.
-    """
-    ds = ImageEmbeddingDataset(
-        tar_url,
-        embeddings_url,
-        shard_width=shard_width,
-        shuffle_shards=shuffle_shards,
-        resample=resample_shards,
-        handler=handler
-    )
-    if shuffle_num is not None and shuffle_num > 0:
-        ds.shuffle(1000)
-    return wds.WebLoader(
-        ds,
-        num_workers=num_workers,
-        batch_size=batch_size,
-        prefetch_factor=2,  # This might be good to have high so the next npy file is prefetched
-        pin_memory=True,
-        shuffle=False
-    )

dalle2_pytorch/optimizer.py

@@ -1,30 +0,0 @@
-from torch.optim import AdamW, Adam
-
-def separate_weight_decayable_params(params):
-    no_wd_params = set([param for param in params if param.ndim < 2])
-    wd_params = set(params) - no_wd_params
-    return wd_params, no_wd_params
-
-def get_optimizer(
-    params,
-    lr = 2e-5,
-    wd = 1e-2,
-    betas = (0.9, 0.999),
-    eps = 1e-8,
-    filter_by_requires_grad = False
-):
-    if filter_by_requires_grad:
-        params = list(filter(lambda t: t.requires_grad, params))
-
-    if wd == 0:
-        return Adam(params, lr = lr, betas = betas, eps = eps)
-
-    params = set(params)
-    wd_params, no_wd_params = separate_weight_decayable_params(params)
-
-    param_groups = [
-        {'params': list(wd_params)},
-        {'params': list(no_wd_params), 'weight_decay': 0},
-    ]
-
-    return AdamW(param_groups, lr = lr, weight_decay = wd, betas = betas, eps = eps)

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,440 +0,0 @@
-import time
-import copy
-from math import ceil
-from functools import partial, wraps
-from collections.abc import Iterable
-
-import torch
-from torch import nn
-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):
-    return val is not None
-
-def default(val, d):
-    return val if exists(val) else d
-
-def cast_tuple(val, length = 1):
-    return val if isinstance(val, tuple) else ((val,) * length)
-
-def pick_and_pop(keys, d):
-    values = list(map(lambda key: d.pop(key), keys))
-    return dict(zip(keys, values))
-
-def group_dict_by_key(cond, d):
-    return_val = [dict(),dict()]
-    for key in d.keys():
-        match = bool(cond(key))
-        ind = int(not match)
-        return_val[ind][key] = d[key]
-    return (*return_val,)
-
-def string_begins_with(prefix, str):
-    return str.startswith(prefix)
-
-def group_by_key_prefix(prefix, d):
-    return group_dict_by_key(partial(string_begins_with, prefix), d)
-
-def groupby_prefix_and_trim(prefix, d):
-    kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, 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):
-    accum = []
-    for ind in range(ceil(len(it) / split_size)):
-        start_index = ind * split_size
-        accum.append(it[start_index: (start_index + split_size)])
-    return accum
-
-def split(t, split_size = None):
-    if not exists(split_size):
-        return t
-
-    if isinstance(t, torch.Tensor):
-        return t.split(split_size, dim = 0)
-
-    if isinstance(t, Iterable):
-        return split_iterable(t, split_size)
-
-    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)
-    split_size = default(split_size, batch_size)
-    num_chunks = ceil(batch_size / split_size)
-
-    dict_len = len(kwargs)
-    dict_keys = kwargs.keys()
-    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]
-    chunk_sizes = tuple(map(len, split_all_args[0]))
-
-    for (chunk_size, *chunked_all_args) in tuple(zip(chunk_sizes, *split_all_args)):
-        chunked_args, chunked_kwargs_values = chunked_all_args[:split_kwargs_index], chunked_all_args[split_kwargs_index:]
-        chunked_kwargs = dict(tuple(zip(dict_keys, chunked_kwargs_values)))
-        chunk_size_frac = chunk_size / batch_size
-        yield chunk_size_frac, (chunked_args, chunked_kwargs)
-
-# print helpers
-
-def print_ribbon(s, symbol = '=', repeat = 40):
-    flank = symbol * repeat
-    return f'{flank} {s} {flank}'
-
-# saving and loading functions
-
-# for diffusion prior
-
-def load_diffusion_model(dprior_path, device):
-    dprior_path = Path(dprior_path)
-    assert dprior_path.exists(), 'Dprior model file does not exist'
-    loaded_obj = torch.load(str(dprior_path), map_location='cpu')
-
-    # Get hyperparameters of loaded model
-    dpn_config = loaded_obj['hparams']['diffusion_prior_network']
-    dp_config = loaded_obj['hparams']['diffusion_prior']
-    image_embed_dim = loaded_obj['image_embed_dim']['image_embed_dim']
-
-    # Create DiffusionPriorNetwork and DiffusionPrior with loaded hyperparameters
-
-    # DiffusionPriorNetwork
-    prior_network = DiffusionPriorNetwork( dim = image_embed_dim, **dpn_config).to(device)
-
-    # DiffusionPrior with text embeddings and image embeddings pre-computed
-    diffusion_prior = DiffusionPrior(net = prior_network, **dp_config, image_embed_dim = image_embed_dim).to(device)
-
-    # Load state dict from saved model
-    diffusion_prior.load_state_dict(loaded_obj['model'])
-
-    return diffusion_prior, loaded_obj
-
-def save_diffusion_model(save_path, model, optimizer, scaler, config, image_embed_dim):
-    # Saving State Dict
-    print_ribbon('Saving checkpoint')
-
-    state_dict = dict(model=model.state_dict(),
-                      optimizer=optimizer.state_dict(),
-                      scaler=scaler.state_dict(),
-                      hparams = config,
-                      image_embed_dim = {"image_embed_dim":image_embed_dim})
-    torch.save(state_dict, save_path+'/'+str(time.time())+'_saved_model.pth')
-
-# exponential moving average wrapper
-
-class EMA(nn.Module):
-    def __init__(
-        self,
-        model,
-        beta = 0.9999,
-        update_after_step = 1000,
-        update_every = 10,
-    ):
-        super().__init__()
-        self.beta = beta
-        self.online_model = model
-        self.ema_model = copy.deepcopy(model)
-
-        self.update_after_step = update_after_step # only start EMA after this step number, starting at 0
-        self.update_every = update_every
-
-        self.register_buffer('initted', torch.Tensor([False]))
-        self.register_buffer('step', torch.tensor([0.]))
-
-    def restore_ema_model_device(self):
-        device = self.initted.device
-        self.ema_model.to(device)
-
-    def update(self):
-        self.step += 1
-
-        if self.step <= self.update_after_step or (self.step % self.update_every) != 0:
-            return
-
-        if not self.initted:
-            self.ema_model.state_dict(self.online_model.state_dict())
-            self.initted.data.copy_(torch.Tensor([True]))
-
-        self.update_moving_average(self.ema_model, self.online_model)
-
-    def update_moving_average(self, ma_model, current_model):
-        def calculate_ema(beta, old, new):
-            if not exists(old):
-                return new
-            return old * beta + (1 - beta) * new
-
-        for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
-            old_weight, up_weight = ma_params.data, current_params.data
-            ma_params.data = calculate_ema(self.beta, old_weight, up_weight)
-
-        for current_buffer, ma_buffer in zip(current_model.buffers(), ma_model.buffers()):
-            new_buffer_value = calculate_ema(self.beta, ma_buffer, current_buffer)
-            ma_buffer.copy_(new_buffer_value)
-
-    def __call__(self, *args, **kwargs):
-        return self.ema_model(*args, **kwargs)
-
-# diffusion prior trainer
-
-class DiffusionPriorTrainer(nn.Module):
-    def __init__(
-        self,
-        diffusion_prior,
-        use_ema = True,
-        lr = 3e-4,
-        wd = 1e-2,
-        eps = 1e-6,
-        max_grad_norm = None,
-        amp = False,
-        **kwargs
-    ):
-        super().__init__()
-        assert isinstance(diffusion_prior, DiffusionPrior)
-        ema_kwargs, kwargs = groupby_prefix_and_trim('ema_', kwargs)
-
-        self.diffusion_prior = diffusion_prior
-
-        # exponential moving average
-
-        self.use_ema = use_ema
-        if self.use_ema:
-            self.ema_diffusion_prior = EMA(diffusion_prior, **ema_kwargs)
-
-        # optimizer and mixed precision stuff
-
-        self.amp = amp
-
-        self.scaler = GradScaler(enabled = amp)
-
-        self.optimizer = get_optimizer(
-            diffusion_prior.parameters(),
-            lr = lr,
-            wd = wd,
-            eps = eps,
-            **kwargs
-        )
-
-        # gradient clipping if needed
-
-        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)
-            nn.utils.clip_grad_norm_(self.diffusion_prior.parameters(), self.max_grad_norm)
-
-        self.scaler.step(self.optimizer)
-        self.scaler.update()
-        self.optimizer.zero_grad()
-
-        if self.use_ema:
-            self.ema_diffusion_prior.update()
-
-        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)
-
-    @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,
-        *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):
-            with autocast(enabled = self.amp):
-                loss = self.diffusion_prior(*chunked_args, **chunked_kwargs)
-                loss = loss * chunk_size_frac
-
-            total_loss += loss.item()
-
-            if self.training:
-                self.scaler.scale(loss).backward()
-
-        return total_loss
-
-# decoder trainer
-
-class DecoderTrainer(nn.Module):
-    def __init__(
-        self,
-        decoder,
-        use_ema = True,
-        lr = 2e-5,
-        wd = 1e-2,
-        eps = 1e-8,
-        max_grad_norm = None,
-        amp = False,
-        **kwargs
-    ):
-        super().__init__()
-        assert isinstance(decoder, Decoder)
-        ema_kwargs, kwargs = groupby_prefix_and_trim('ema_', kwargs)
-
-        self.decoder = decoder
-        self.num_unets = len(self.decoder.unets)
-
-        self.use_ema = use_ema
-        self.ema_unets = nn.ModuleList([])
-
-        self.amp = amp
-
-        # be able to finely customize learning rate, weight decay
-        # per unet
-
-        lr, wd, eps = map(partial(cast_tuple, length = self.num_unets), (lr, wd, eps))
-
-        for ind, (unet, unet_lr, unet_wd, unet_eps) in enumerate(zip(self.decoder.unets, lr, wd, eps)):
-            optimizer = get_optimizer(
-                unet.parameters(),
-                lr = unet_lr,
-                wd = unet_wd,
-                eps = unet_eps,
-                **kwargs
-            )
-
-            setattr(self, f'optim{ind}', optimizer) # cannot use pytorch ModuleList for some reason with optimizers
-
-            if self.use_ema:
-                self.ema_unets.append(EMA(unet, **ema_kwargs))
-
-            scaler = GradScaler(enabled = amp)
-            setattr(self, f'scaler{ind}', scaler)
-
-        # gradient clipping if needed
-
-        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])
-
-    def scale(self, loss, *, unet_number):
-        assert 1 <= unet_number <= self.num_unets
-        index = unet_number - 1
-        scaler = getattr(self, f'scaler{index}')
-        return scaler.scale(loss)
-
-    def update(self, unet_number):
-        assert 1 <= unet_number <= self.num_unets
-        index = unet_number - 1
-        unet = self.decoder.unets[index]
-
-        optimizer = getattr(self, f'optim{index}')
-        scaler = getattr(self, f'scaler{index}')
-
-        if exists(self.max_grad_norm):
-            scaler.unscale_(optimizer)
-            nn.utils.clip_grad_norm_(unet.parameters(), self.max_grad_norm)
-
-        scaler.step(optimizer)
-        scaler.update()
-        optimizer.zero_grad()
-
-        if self.use_ema:
-            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
-            self.decoder.unets = self.unets                  # swap in exponential moving averaged unets for sampling
-
-        output = self.decoder.sample(*args, **kwargs)
-
-        if self.use_ema:
-            self.decoder.unets = trainable_unets             # restore original training unets
-
-        # cast the ema_model unets back to original device
-        for ema in self.ema_unets:
-            ema.restore_ema_model_device()
-
-        return output
-
-    @cast_torch_tensor
-    def forward(
-        self,
-        *args,
-        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):
-            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()
-
-            if self.training:
-                self.scale(loss, unet_number = unet_number).backward()
-
-        return total_loss

dalle2_pytorch/train_vqgan_vae.py

@@ -1,277 +0,0 @@
-from math import sqrt
-import copy
-from random import choice
-from pathlib import Path
-from shutil import rmtree
-from PIL import Image
-
-import torch
-from torch import nn
-from torch.cuda.amp import autocast, GradScaler
-from torch.utils.data import Dataset, DataLoader, random_split
-
-import torchvision.transforms as T
-from torchvision.datasets import ImageFolder
-from torchvision.utils import make_grid, save_image
-
-from einops import rearrange
-
-from dalle2_pytorch.train import EMA
-from dalle2_pytorch.vqgan_vae import VQGanVAE
-from dalle2_pytorch.optimizer import get_optimizer
-
-# helpers
-
-def exists(val):
-    return val is not None
-
-def noop(*args, **kwargs):
-    pass
-
-def cycle(dl):
-    while True:
-        for data in dl:
-            yield data
-
-def cast_tuple(t):
-    return t if isinstance(t, (tuple, list)) else (t,)
-
-def yes_or_no(question):
-    answer = input(f'{question} (y/n) ')
-    return answer.lower() in ('yes', 'y')
-
-def accum_log(log, new_logs):
-    for key, new_value in new_logs.items():
-        old_value = log.get(key, 0.)
-        log[key] = old_value + new_value
-    return log
-
-# classes
-
-class ImageDataset(Dataset):
-    def __init__(
-        self,
-        folder,
-        image_size,
-        exts = ['jpg', 'jpeg', 'png']
-    ):
-        super().__init__()
-        self.folder = folder
-        self.image_size = image_size
-        self.paths = [p for ext in exts for p in Path(f'{folder}').glob(f'**/*.{ext}')]
-
-        print(f'{len(self.paths)} training samples found at {folder}')
-
-        self.transform = T.Compose([
-            T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
-            T.Resize(image_size),
-            T.RandomHorizontalFlip(),
-            T.CenterCrop(image_size),
-            T.ToTensor()
-        ])
-
-    def __len__(self):
-        return len(self.paths)
-
-    def __getitem__(self, index):
-        path = self.paths[index]
-        img = Image.open(path)
-        return self.transform(img)
-
-# main trainer class
-
-class VQGanVAETrainer(nn.Module):
-    def __init__(
-        self,
-        vae,
-        *,
-        num_train_steps,
-        lr,
-        batch_size,
-        folder,
-        grad_accum_every,
-        wd = 0.,
-        save_results_every = 100,
-        save_model_every = 1000,
-        results_folder = './results',
-        valid_frac = 0.05,
-        random_split_seed = 42,
-        ema_beta = 0.995,
-        ema_update_after_step = 2000,
-        ema_update_every = 10,
-        apply_grad_penalty_every = 4,
-        amp = False
-    ):
-        super().__init__()
-        assert isinstance(vae, VQGanVAE), 'vae must be instance of VQGanVAE'
-        image_size = vae.image_size
-
-        self.vae = vae
-        self.ema_vae = EMA(vae, update_after_step = ema_update_after_step, update_every = ema_update_every)
-
-        self.register_buffer('steps', torch.Tensor([0]))
-
-        self.num_train_steps = num_train_steps
-        self.batch_size = batch_size
-        self.grad_accum_every = grad_accum_every
-
-        all_parameters = set(vae.parameters())
-        discr_parameters = set(vae.discr.parameters())
-        vae_parameters = all_parameters - discr_parameters
-
-        self.optim = get_optimizer(vae_parameters, lr = lr, wd = wd)
-        self.discr_optim = get_optimizer(discr_parameters, lr = lr, wd = wd)
-
-        self.amp = amp
-        self.scaler = GradScaler(enabled = amp)
-        self.discr_scaler = GradScaler(enabled = amp)
-
-        # create dataset
-
-        self.ds = ImageDataset(folder, image_size = image_size)
-
-        # split for validation
-
-        if valid_frac > 0:
-            train_size = int((1 - valid_frac) * len(self.ds))
-            valid_size = len(self.ds) - train_size
-            self.ds, self.valid_ds = random_split(self.ds, [train_size, valid_size], generator = torch.Generator().manual_seed(random_split_seed))
-            print(f'training with dataset of {len(self.ds)} samples and validating with randomly splitted {len(self.valid_ds)} samples')
-        else:
-            self.valid_ds = self.ds
-            print(f'training with shared training and valid dataset of {len(self.ds)} samples')
-
-        # dataloader
-
-        self.dl = cycle(DataLoader(
-            self.ds,
-            batch_size = batch_size,
-            shuffle = True
-        ))
-
-        self.valid_dl = cycle(DataLoader(
-            self.valid_ds,
-            batch_size = batch_size,
-            shuffle = True
-        ))
-
-        self.save_model_every = save_model_every
-        self.save_results_every = save_results_every
-
-        self.apply_grad_penalty_every = apply_grad_penalty_every
-
-        self.results_folder = Path(results_folder)
-
-        if len([*self.results_folder.glob('**/*')]) > 0 and yes_or_no('do you want to clear previous experiment checkpoints and results?'):
-            rmtree(str(self.results_folder))
-
-        self.results_folder.mkdir(parents = True, exist_ok = True)
-
-    def train_step(self):
-        device = next(self.vae.parameters()).device
-        steps = int(self.steps.item())
-        apply_grad_penalty = not (steps % self.apply_grad_penalty_every)
-
-        self.vae.train()
-
-        # logs
-
-        logs = {}
-
-        # update vae (generator)
-
-        for _ in range(self.grad_accum_every):
-            img = next(self.dl)
-            img = img.to(device)
-
-            with autocast(enabled = self.amp):
-                loss = self.vae(
-                    img,
-                    return_loss = True,
-                    apply_grad_penalty = apply_grad_penalty
-                )
-
-
-                self.scaler.scale(loss / self.grad_accum_every).backward()
-
-            accum_log(logs, {'loss': loss.item() / self.grad_accum_every})
-
-        self.scaler.step(self.optim)
-        self.scaler.update()
-        self.optim.zero_grad()
-
-        # update discriminator
-
-        if exists(self.vae.discr):
-            discr_loss = 0
-            for _ in range(self.grad_accum_every):
-                img = next(self.dl)
-                img = img.to(device)
-
-                with autocast(enabled = self.amp):
-                    loss = self.vae(img, return_discr_loss = True)
-
-                    self.discr_scaler.scale(loss / self.grad_accum_every).backward()
-
-                accum_log(logs, {'discr_loss': loss.item() / self.grad_accum_every})
-
-            self.discr_scaler.step(self.discr_optim)
-            self.discr_scaler.update()
-            self.discr_optim.zero_grad()
-
-            # log
-
-            print(f"{steps}: vae loss: {logs['loss']} - discr loss: {logs['discr_loss']}")
-
-        # update exponential moving averaged generator
-
-        self.ema_vae.update()
-
-        # sample results every so often
-
-        if not (steps % self.save_results_every):
-            for model, filename in ((self.ema_vae.ema_model, f'{steps}.ema'), (self.vae, str(steps))):
-                model.eval()
-
-                imgs = next(self.dl)
-                imgs = imgs.to(device)
-
-                recons = model(imgs)
-                nrows = int(sqrt(self.batch_size))
-
-                imgs_and_recons = torch.stack((imgs, recons), dim = 0)
-                imgs_and_recons = rearrange(imgs_and_recons, 'r b ... -> (b r) ...')
-
-                imgs_and_recons = imgs_and_recons.detach().cpu().float().clamp(0., 1.)
-                grid = make_grid(imgs_and_recons, nrow = 2, normalize = True, value_range = (0, 1))
-
-                logs['reconstructions'] = grid
-
-                save_image(grid, str(self.results_folder / f'{filename}.png'))
-
-            print(f'{steps}: saving to {str(self.results_folder)}')
-
-        # save model every so often
-
-        if not (steps % self.save_model_every):
-            state_dict = self.vae.state_dict()
-            model_path = str(self.results_folder / f'vae.{steps}.pt')
-            torch.save(state_dict, model_path)
-
-            ema_state_dict = self.ema_vae.state_dict()
-            model_path = str(self.results_folder / f'vae.{steps}.ema.pt')
-            torch.save(ema_state_dict, model_path)
-
-            print(f'{steps}: saving model to {str(self.results_folder)}')
-
-        self.steps += 1
-        return logs
-
-    def train(self, log_fn = noop):
-        device = next(self.vae.parameters()).device
-
-        while self.steps < self.num_train_steps:
-            logs = self.train_step()
-            log_fn(logs)
-
-        print('training complete')

dalle2_pytorch/vqgan_vae.py

@@ -1,765 +0,0 @@
-import copy
-import math
-from math import sqrt
-from functools import partial, wraps
-
-from vector_quantize_pytorch import VectorQuantize as VQ
-
-import torch
-from torch import nn, einsum
-import torch.nn.functional as F
-from torch.autograd import grad as torch_grad
-import torchvision
-
-from einops import rearrange, reduce, repeat
-from einops_exts import rearrange_many
-from einops.layers.torch import Rearrange
-
-# constants
-
-MList = nn.ModuleList
-
-# helper functions
-
-def exists(val):
-    return val is not None
-
-def default(val, d):
-    return val if exists(val) else d
-
-# decorators
-
-def eval_decorator(fn):
-    def inner(model, *args, **kwargs):
-        was_training = model.training
-        model.eval()
-        out = fn(model, *args, **kwargs)
-        model.train(was_training)
-        return out
-    return inner
-
-def remove_vgg(fn):
-    @wraps(fn)
-    def inner(self, *args, **kwargs):
-        has_vgg = hasattr(self, 'vgg')
-        if has_vgg:
-            vgg = self.vgg
-            delattr(self, 'vgg')
-
-        out = fn(self, *args, **kwargs)
-
-        if has_vgg:
-            self.vgg = vgg
-
-        return out
-    return inner
-
-# keyword argument helpers
-
-def pick_and_pop(keys, d):
-    values = list(map(lambda key: d.pop(key), keys))
-    return dict(zip(keys, values))
-
-def group_dict_by_key(cond, d):
-    return_val = [dict(),dict()]
-    for key in d.keys():
-        match = bool(cond(key))
-        ind = int(not match)
-        return_val[ind][key] = d[key]
-    return (*return_val,)
-
-def string_begins_with(prefix, str):
-    return str.startswith(prefix)
-
-def group_by_key_prefix(prefix, d):
-    return group_dict_by_key(partial(string_begins_with, prefix), d)
-
-def groupby_prefix_and_trim(prefix, d):
-    kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, 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
-
-# tensor helper functions
-
-def log(t, eps = 1e-10):
-    return torch.log(t + eps)
-
-def gradient_penalty(images, output, weight = 10):
-    batch_size = images.shape[0]
-    gradients = torch_grad(outputs = output, inputs = images,
-                           grad_outputs = torch.ones(output.size(), device = images.device),
-                           create_graph = True, retain_graph = True, only_inputs = True)[0]
-
-    gradients = rearrange(gradients, 'b ... -> b (...)')
-    return weight * ((gradients.norm(2, dim = 1) - 1) ** 2).mean()
-
-def l2norm(t):
-    return F.normalize(t, dim = -1)
-
-def leaky_relu(p = 0.1):
-    return nn.LeakyReLU(0.1)
-
-def stable_softmax(t, dim = -1, alpha = 32 ** 2):
-    t = t / alpha
-    t = t - torch.amax(t, dim = dim, keepdim = True).detach()
-    return (t * alpha).softmax(dim = dim)
-
-def safe_div(numer, denom, eps = 1e-8):
-    return numer / (denom + eps)
-
-# gan losses
-
-def hinge_discr_loss(fake, real):
-    return (F.relu(1 + fake) + F.relu(1 - real)).mean()
-
-def hinge_gen_loss(fake):
-    return -fake.mean()
-
-def bce_discr_loss(fake, real):
-    return (-log(1 - torch.sigmoid(fake)) - log(torch.sigmoid(real))).mean()
-
-def bce_gen_loss(fake):
-    return -log(torch.sigmoid(fake)).mean()
-
-def grad_layer_wrt_loss(loss, layer):
-    return torch_grad(
-        outputs = loss,
-        inputs = layer,
-        grad_outputs = torch.ones_like(loss),
-        retain_graph = True
-    )[0].detach()
-
-# vqgan vae
-
-class LayerNormChan(nn.Module):
-    def __init__(
-        self,
-        dim,
-        eps = 1e-5
-    ):
-        super().__init__()
-        self.eps = eps
-        self.gamma = nn.Parameter(torch.ones(1, dim, 1, 1))
-
-    def forward(self, x):
-        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
-        mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) / (var + self.eps).sqrt() * self.gamma
-
-# discriminator
-
-class Discriminator(nn.Module):
-    def __init__(
-        self,
-        dims,
-        channels = 3,
-        groups = 16,
-        init_kernel_size = 5
-    ):
-        super().__init__()
-        dim_pairs = zip(dims[:-1], dims[1:])
-
-        self.layers = MList([nn.Sequential(nn.Conv2d(channels, dims[0], init_kernel_size, padding = init_kernel_size // 2), leaky_relu())])
-
-        for dim_in, dim_out in dim_pairs:
-            self.layers.append(nn.Sequential(
-                nn.Conv2d(dim_in, dim_out, 4, stride = 2, padding = 1),
-                nn.GroupNorm(groups, dim_out),
-                leaky_relu()
-            ))
-
-        dim = dims[-1]
-        self.to_logits = nn.Sequential( # return 5 x 5, for PatchGAN-esque training
-            nn.Conv2d(dim, dim, 1),
-            leaky_relu(),
-            nn.Conv2d(dim, 1, 4)
-        )
-
-    def forward(self, x):
-        for net in self.layers:
-            x = net(x)
-
-        return self.to_logits(x)
-
-# positional encoding
-
-class ContinuousPositionBias(nn.Module):
-    """ from https://arxiv.org/abs/2111.09883 """
-
-    def __init__(self, *, dim, heads, layers = 2):
-        super().__init__()
-        self.net = MList([])
-        self.net.append(nn.Sequential(nn.Linear(2, dim), leaky_relu()))
-
-        for _ in range(layers - 1):
-            self.net.append(nn.Sequential(nn.Linear(dim, dim), leaky_relu()))
-
-        self.net.append(nn.Linear(dim, heads))
-        self.register_buffer('rel_pos', None, persistent = False)
-
-    def forward(self, x):
-        n, device = x.shape[-1], x.device
-        fmap_size = int(sqrt(n))
-
-        if not exists(self.rel_pos):
-            pos = torch.arange(fmap_size, device = device)
-            grid = torch.stack(torch.meshgrid(pos, pos, indexing = 'ij'))
-            grid = rearrange(grid, 'c i j -> (i j) c')
-            rel_pos = rearrange(grid, 'i c -> i 1 c') - rearrange(grid, 'j c -> 1 j c')
-            rel_pos = torch.sign(rel_pos) * torch.log(rel_pos.abs() + 1)
-            self.register_buffer('rel_pos', rel_pos, persistent = False)
-
-        rel_pos = self.rel_pos.float()
-
-        for layer in self.net:
-            rel_pos = layer(rel_pos)
-
-        bias = rearrange(rel_pos, 'i j h -> h i j')
-        return x + bias
-
-# resnet encoder / decoder
-
-class ResnetEncDec(nn.Module):
-    def __init__(
-        self,
-        dim,
-        *,
-        channels = 3,
-        layers = 4,
-        layer_mults = None,
-        num_resnet_blocks = 1,
-        resnet_groups = 16,
-        first_conv_kernel_size = 5,
-        use_attn = True,
-        attn_dim_head = 64,
-        attn_heads = 8,
-        attn_dropout = 0.,
-    ):
-        super().__init__()
-        assert dim % resnet_groups == 0, f'dimension {dim} must be divisible by {resnet_groups} (groups for the groupnorm)'
-
-        self.layers = layers
-
-        self.encoders = MList([])
-        self.decoders = MList([])
-
-        layer_mults = default(layer_mults, list(map(lambda t: 2 ** t, range(layers))))
-        assert len(layer_mults) == layers, 'layer multipliers must be equal to designated number of layers'
-
-        layer_dims = [dim * mult for mult in layer_mults]
-        dims = (dim, *layer_dims)
-
-        self.encoded_dim = dims[-1]
-
-        dim_pairs = zip(dims[:-1], dims[1:])
-
-        append = lambda arr, t: arr.append(t)
-        prepend = lambda arr, t: arr.insert(0, t)
-
-        if not isinstance(num_resnet_blocks, tuple):
-            num_resnet_blocks = (*((0,) * (layers - 1)), num_resnet_blocks)
-
-        if not isinstance(use_attn, tuple):
-            use_attn = (*((False,) * (layers - 1)), use_attn)
-
-        assert len(num_resnet_blocks) == layers, 'number of resnet blocks config must be equal to number of layers'
-        assert len(use_attn) == layers
-
-        for layer_index, (dim_in, dim_out), layer_num_resnet_blocks, layer_use_attn in zip(range(layers), dim_pairs, num_resnet_blocks, use_attn):
-            append(self.encoders, nn.Sequential(nn.Conv2d(dim_in, dim_out, 4, stride = 2, padding = 1), leaky_relu()))
-            prepend(self.decoders, nn.Sequential(nn.ConvTranspose2d(dim_out, dim_in, 4, 2, 1), leaky_relu()))
-
-            if layer_use_attn:
-                prepend(self.decoders, VQGanAttention(dim = dim_out, heads = attn_heads, dim_head = attn_dim_head, dropout = attn_dropout))
-
-            for _ in range(layer_num_resnet_blocks):
-                append(self.encoders, ResBlock(dim_out, groups = resnet_groups))
-                prepend(self.decoders, GLUResBlock(dim_out, groups = resnet_groups))
-
-            if layer_use_attn:
-                append(self.encoders, VQGanAttention(dim = dim_out, heads = attn_heads, dim_head = attn_dim_head, dropout = attn_dropout))
-
-        prepend(self.encoders, nn.Conv2d(channels, dim, first_conv_kernel_size, padding = first_conv_kernel_size // 2))
-        append(self.decoders, nn.Conv2d(dim, channels, 1))
-
-    def get_encoded_fmap_size(self, image_size):
-        return image_size // (2 ** self.layers)
-
-    @property
-    def last_dec_layer(self):
-        return self.decoders[-1].weight
-
-    def encode(self, x):
-        for enc in self.encoders:
-            x = enc(x)
-        return x
-
-    def decode(self, x):
-        for dec in self.decoders:
-            x = dec(x)
-        return x
-
-class GLUResBlock(nn.Module):
-    def __init__(self, chan, groups = 16):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Conv2d(chan, chan * 2, 3, padding = 1),
-            nn.GLU(dim = 1),
-            nn.GroupNorm(groups, chan),
-            nn.Conv2d(chan, chan * 2, 3, padding = 1),
-            nn.GLU(dim = 1),
-            nn.GroupNorm(groups, chan),
-            nn.Conv2d(chan, chan, 1)
-        )
-
-    def forward(self, x):
-        return self.net(x) + x
-
-class ResBlock(nn.Module):
-    def __init__(self, chan, groups = 16):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Conv2d(chan, chan, 3, padding = 1),
-            nn.GroupNorm(groups, chan),
-            leaky_relu(),
-            nn.Conv2d(chan, chan, 3, padding = 1),
-            nn.GroupNorm(groups, chan),
-            leaky_relu(),
-            nn.Conv2d(chan, chan, 1)
-        )
-
-    def forward(self, x):
-        return self.net(x) + x
-
-# vqgan attention layer
-
-class VQGanAttention(nn.Module):
-    def __init__(
-        self,
-        *,
-        dim,
-        dim_head = 64,
-        heads = 8,
-        dropout = 0.
-    ):
-        super().__init__()
-        self.heads = heads
-        self.scale = dim_head ** -0.5
-        inner_dim = heads * dim_head
-
-        self.dropout = nn.Dropout(dropout)
-        self.pre_norm = LayerNormChan(dim)
-
-        self.cpb = ContinuousPositionBias(dim = dim // 4, heads = heads)
-        self.to_qkv = nn.Conv2d(dim, inner_dim * 3, 1, bias = False)
-        self.to_out = nn.Conv2d(inner_dim, dim, 1, bias = False)
-
-    def forward(self, x):
-        h = self.heads
-        height, width, residual = *x.shape[-2:], x.clone()
-
-        x = self.pre_norm(x)
-
-        q, k, v = self.to_qkv(x).chunk(3, dim = 1)
-
-        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h = h), (q, k, v))
-
-        sim = einsum('b h c i, b h c j -> b h i j', q, k) * self.scale
-
-        sim = self.cpb(sim)
-
-        attn = stable_softmax(sim, dim = -1)
-        attn = self.dropout(attn)
-
-        out = einsum('b h i j, b h c j -> b h c i', attn, v)
-        out = rearrange(out, 'b h c (x y) -> b (h c) x y', x = height, y = width)
-        out = self.to_out(out)
-
-        return out + residual
-
-# ViT encoder / decoder
-
-class RearrangeImage(nn.Module):
-    def forward(self, x):
-        n = x.shape[1]
-        w = h = int(sqrt(n))
-        return rearrange(x, 'b (h w) ... -> b h w ...', h = h, w = w)
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        dim,
-        *,
-        heads = 8,
-        dim_head = 32
-    ):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.heads = heads
-        self.scale = dim_head ** -0.5
-        inner_dim = dim_head * heads
-
-        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
-        self.to_out = nn.Linear(inner_dim, dim)
-
-    def forward(self, x):
-        h = self.heads
-
-        x = self.norm(x)
-
-        q, k, v = self.to_qkv(x).chunk(3, dim = -1)
-        q, k, v = rearrange_many((q, k, v), 'b n (h d) -> b h n d', h = h)
-
-        q = q * self.scale
-        sim = einsum('b h i d, b h j d -> b h i j', q, k)
-
-        sim = sim - sim.amax(dim = -1, keepdim = True).detach()
-        attn = sim.softmax(dim = -1)
-
-        out = einsum('b h i j, b h j d -> b h i d', attn, v)
-
-        out = rearrange(out, 'b h n d -> b n (h d)')
-        return self.to_out(out)
-
-def FeedForward(dim, mult = 4):
-    return nn.Sequential(
-        nn.LayerNorm(dim),
-        nn.Linear(dim, dim * mult, bias = False),
-        nn.GELU(),
-        nn.Linear(dim * mult, dim, bias = False)
-    )
-
-class Transformer(nn.Module):
-    def __init__(
-        self,
-        dim,
-        *,
-        layers,
-        dim_head = 32,
-        heads = 8,
-        ff_mult = 4
-    ):
-        super().__init__()
-        self.layers = nn.ModuleList([])
-        for _ in range(layers):
-            self.layers.append(nn.ModuleList([
-                Attention(dim = dim, dim_head = dim_head, heads = heads),
-                FeedForward(dim = dim, mult = ff_mult)
-            ]))
-
-        self.norm = nn.LayerNorm(dim)
-
-    def forward(self, x):
-        for attn, ff in self.layers:
-            x = attn(x) + x
-            x = ff(x) + x
-
-        return self.norm(x)
-
-class ViTEncDec(nn.Module):
-    def __init__(
-        self,
-        dim,
-        channels = 3,
-        layers = 4,
-        patch_size = 8,
-        dim_head = 32,
-        heads = 8,
-        ff_mult = 4
-    ):
-        super().__init__()
-        self.encoded_dim = dim
-        self.patch_size = patch_size
-
-        input_dim = channels * (patch_size ** 2)
-
-        self.encoder = nn.Sequential(
-            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
-            nn.Linear(input_dim, dim),
-            Transformer(
-                dim = dim,
-                dim_head = dim_head,
-                heads = heads,
-                ff_mult = ff_mult,
-                layers = layers
-            ),
-            RearrangeImage(),
-            Rearrange('b h w c -> b c h w')
-        )
-
-        self.decoder = nn.Sequential(
-            Rearrange('b c h w -> b (h w) c'),
-            Transformer(
-                dim = dim,
-                dim_head = dim_head,
-                heads = heads,
-                ff_mult = ff_mult,
-                layers = layers
-            ),
-            nn.Sequential(
-                nn.Linear(dim, dim * 4, bias = False),
-                nn.Tanh(),
-                nn.Linear(dim * 4, input_dim, bias = False),
-            ),
-            RearrangeImage(),
-            Rearrange('b h w (p1 p2 c) -> b c (h p1) (w p2)', p1 = patch_size, p2 = patch_size)
-        )
-
-    def get_encoded_fmap_size(self, image_size):
-        return image_size // self.patch_size
-
-    @property
-    def last_dec_layer(self):
-        return self.decoder[-3][-1].weight
-
-    def encode(self, x):
-        return self.encoder(x)
-
-    def decode(self, x):
-        return self.decoder(x)
-
-# main vqgan-vae classes
-
-class NullVQGanVAE(nn.Module):
-    def __init__(
-        self,
-        *,
-        channels
-    ):
-        super().__init__()
-        self.encoded_dim = channels
-        self.layers = 0
-
-    def get_encoded_fmap_size(self, size):
-        return size
-
-    def copy_for_eval(self):
-        return self
-
-    def encode(self, x):
-        return x
-
-    def decode(self, x):
-        return x
-
-class VQGanVAE(nn.Module):
-    def __init__(
-        self,
-        *,
-        dim,
-        image_size,
-        channels = 3,
-        layers = 4,
-        l2_recon_loss = False,
-        use_hinge_loss = True,
-        vgg = None,
-        vq_codebook_dim = 256,
-        vq_codebook_size = 512,
-        vq_decay = 0.8,
-        vq_commitment_weight = 1.,
-        vq_kmeans_init = True,
-        vq_use_cosine_sim = True,
-        use_vgg_and_gan = True,
-        vae_type = 'resnet',
-        discr_layers = 4,
-        **kwargs
-    ):
-        super().__init__()
-        vq_kwargs, kwargs = groupby_prefix_and_trim('vq_', kwargs)
-        encdec_kwargs, kwargs = groupby_prefix_and_trim('encdec_', kwargs)
-
-        self.image_size = image_size
-        self.channels = channels
-        self.codebook_size = vq_codebook_size
-
-        if vae_type == 'resnet':
-            enc_dec_klass = ResnetEncDec
-        elif vae_type == 'vit':
-            enc_dec_klass = ViTEncDec
-        else:
-            raise ValueError(f'{vae_type} not valid')
-
-        self.enc_dec = enc_dec_klass(
-            dim = dim,
-            channels = channels,
-            layers = layers,
-            **encdec_kwargs
-        )
-
-        self.vq = VQ(
-            dim = self.enc_dec.encoded_dim,
-            codebook_dim = vq_codebook_dim,
-            codebook_size = vq_codebook_size,
-            decay = vq_decay,
-            commitment_weight = vq_commitment_weight,
-            accept_image_fmap = True,
-            kmeans_init = vq_kmeans_init,
-            use_cosine_sim = vq_use_cosine_sim,
-            **vq_kwargs
-        )
-
-        # reconstruction loss
-
-        self.recon_loss_fn = F.mse_loss if l2_recon_loss else F.l1_loss
-
-        # turn off GAN and perceptual loss if grayscale
-
-        self.vgg = None
-        self.discr = None
-        self.use_vgg_and_gan = use_vgg_and_gan
-
-        if not use_vgg_and_gan:
-            return
-
-        # preceptual loss
-
-        if exists(vgg):
-            self.vgg = vgg
-        else:
-            self.vgg = torchvision.models.vgg16(pretrained = True)
-            self.vgg.classifier = nn.Sequential(*self.vgg.classifier[:-2])
-
-        # gan related losses
-
-        layer_mults = list(map(lambda t: 2 ** t, range(discr_layers)))
-        layer_dims = [dim * mult for mult in layer_mults]
-        dims = (dim, *layer_dims)
-
-        self.discr = Discriminator(dims = dims, channels = channels)
-
-        self.discr_loss = hinge_discr_loss if use_hinge_loss else bce_discr_loss
-        self.gen_loss = hinge_gen_loss if use_hinge_loss else bce_gen_loss
-
-    @property
-    def encoded_dim(self):
-        return self.enc_dec.encoded_dim
-
-    def get_encoded_fmap_size(self, image_size):
-        return self.enc_dec.get_encoded_fmap_size(image_size)
-
-    def copy_for_eval(self):
-        device = next(self.parameters()).device
-        vae_copy = copy.deepcopy(self.cpu())
-
-        if vae_copy.use_vgg_and_gan:
-            del vae_copy.discr
-            del vae_copy.vgg
-
-        vae_copy.eval()
-        return vae_copy.to(device)
-
-    @remove_vgg
-    def state_dict(self, *args, **kwargs):
-        return super().state_dict(*args, **kwargs)
-
-    @remove_vgg
-    def load_state_dict(self, *args, **kwargs):
-        return super().load_state_dict(*args, **kwargs)
-
-    @property
-    def codebook(self):
-        return self.vq.codebook
-
-    def encode(self, fmap):
-        fmap = self.enc_dec.encode(fmap)
-        return fmap
-
-    def decode(self, fmap, return_indices_and_loss = False):
-        fmap, indices, commit_loss = self.vq(fmap)
-
-        fmap = self.enc_dec.decode(fmap)
-
-        if not return_indices_and_loss:
-            return fmap
-
-        return fmap, indices, commit_loss
-
-    def forward(
-        self,
-        img,
-        return_loss = False,
-        return_discr_loss = False,
-        return_recons = False,
-        add_gradient_penalty = True
-    ):
-        batch, channels, height, width, device = *img.shape, img.device
-        assert height == self.image_size and width == self.image_size, 'height and width of input image must be equal to {self.image_size}'
-        assert channels == self.channels, 'number of channels on image or sketch is not equal to the channels set on this VQGanVAE'
-
-        fmap = self.encode(img)
-
-        fmap, indices, commit_loss = self.decode(fmap, return_indices_and_loss = True)
-
-        if not return_loss and not return_discr_loss:
-            return fmap
-
-        assert return_loss ^ return_discr_loss, 'you should either return autoencoder loss or discriminator loss, but not both'
-
-        # whether to return discriminator loss
-
-        if return_discr_loss:
-            assert exists(self.discr), 'discriminator must exist to train it'
-
-            fmap.detach_()
-            img.requires_grad_()
-
-            fmap_discr_logits, img_discr_logits = map(self.discr, (fmap, img))
-
-            discr_loss = self.discr_loss(fmap_discr_logits, img_discr_logits)
-
-            if add_gradient_penalty:
-                gp = gradient_penalty(img, img_discr_logits)
-                loss = discr_loss + gp
-
-            if return_recons:
-                return loss, fmap
-
-            return loss
-
-        # reconstruction loss
-
-        recon_loss = self.recon_loss_fn(fmap, img)
-
-        # early return if training on grayscale
-
-        if not self.use_vgg_and_gan:
-            if return_recons:
-                return recon_loss, fmap
-
-            return recon_loss
-
-        # perceptual loss
-
-        img_vgg_input = img
-        fmap_vgg_input = fmap
-
-        if img.shape[1] == 1:
-            # handle grayscale for vgg
-            img_vgg_input, fmap_vgg_input = map(lambda t: repeat(t, 'b 1 ... -> b c ...', c = 3), (img_vgg_input, fmap_vgg_input))
-
-        img_vgg_feats = self.vgg(img_vgg_input)
-        recon_vgg_feats = self.vgg(fmap_vgg_input)
-        perceptual_loss = F.mse_loss(img_vgg_feats, recon_vgg_feats)
-
-        # generator loss
-
-        gen_loss = self.gen_loss(self.discr(fmap))
-
-        # calculate adaptive weight
-
-        last_dec_layer = self.enc_dec.last_dec_layer
-
-        norm_grad_wrt_gen_loss = grad_layer_wrt_loss(gen_loss, last_dec_layer).norm(p = 2)
-        norm_grad_wrt_perceptual_loss = grad_layer_wrt_loss(perceptual_loss, last_dec_layer).norm(p = 2)
-
-        adaptive_weight = safe_div(norm_grad_wrt_perceptual_loss, norm_grad_wrt_gen_loss)
-        adaptive_weight.clamp_(max = 1e4)
-
-        # combine losses
-
-        loss = recon_loss + perceptual_loss + commit_loss + adaptive_weight * gen_loss
-
-        if return_recons:
-            return loss, fmap
-
-        return loss

setup.py

@@ -10,12 +10,11 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.2.37',
+  version = '0.0.4',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',
   author_email = 'lucidrains@gmail.com',
-  long_description_content_type = 'text/markdown',
   url = 'https://github.com/lucidrains/dalle2-pytorch',
   keywords = [
     'artificial intelligence',
@@ -24,24 +23,14 @@ setup(
   ],
   install_requires=[
     'click',
-    'clip-anytorch',
-    'coca-pytorch>=0.0.5',
     'einops>=0.4',
     'einops-exts>=0.0.3',
-    'embedding-reader',
-    'kornia>=0.5.4',
-    'numpy',
     'pillow',
-    'resize-right>=0.0.2',
-    'rotary-embedding-torch',
     'torch>=1.10',
     'torchvision',
     'tqdm',
-    'vector-quantize-pytorch',
     'x-clip>=0.4.4',
-    'youtokentome',
-    'webdataset>=0.2.5',
-    'fsspec>=2022.1.0'
+    'youtokentome'
   ],
   classifiers=[
     'Development Status :: 4 - Beta',

train_diffusion_prior.py

@@ -1,346 +0,0 @@
-from pathlib import Path
-import click
-import math
-import time
-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 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():
-        total_loss = 0.
-        total_samples = 0.
-
-        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)):
-
-            emb_images_tensor = torch.tensor(emb_images[0]).to(device)
-            emb_text_tensor = torch.tensor(emb_text[0]).to(device)
-
-            batches = emb_images_tensor.shape[0]
-
-            loss = model(text_embed = emb_text_tensor, image_embed = emb_images_tensor)
-
-            total_loss += loss.item() * batches
-            total_samples += batches
-
-        avg_loss = (total_loss / total_samples)
-        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):
-    diffusion_prior.eval()
-
-    cos = nn.CosineSimilarity(dim=1, eps=1e-6)
-
-    tstart = train_set_size
-    tend = train_set_size+NUM_TEST_EMBEDDINGS
-
-    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_embed / text_embed.norm(dim=1, keepdim=True)
-       test_text_cond = dict(text_embed=text_embed)
-        # prepare image embeddings
-       test_image_embeddings = torch.tensor(embi[0]).to(device)
-       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(
-           (NUM_TEST_EMBEDDINGS, 768), text_cond=test_text_cond)
-       predicted_image_embeddings = predicted_image_embeddings / \
-           predicted_image_embeddings.norm(dim=1, keepdim=True)
-        # predict on the shuffled embeddings
-       predicted_unrelated_embeddings = diffusion_prior.p_sample_loop(
-           (NUM_TEST_EMBEDDINGS, 768), text_cond=test_text_shuffled_cond)
-       predicted_unrelated_embeddings = predicted_unrelated_embeddings / \
-           predicted_unrelated_embeddings.norm(dim=1, keepdim=True)
-        # calculate similarities
-       original_similarity = cos(
-           text_embed, test_image_embeddings).cpu().numpy()
-       predicted_similarity = cos(
-           text_embed, predicted_image_embeddings).cpu().numpy()
-       unrelated_similarity = cos(
-           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),
-            "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)})
-
-@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 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
-
-    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
-    )
-    
-    # diffusion prior 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
-    )
-
-    # 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'])
-
-    # Create save_path if it doesn't exist
-
-    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 ###
-
-    timer = Timer()
-    epochs = num_epochs
-
-    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 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()
-            
-            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)
-
-            # Samples per second
-
-            samples_per_sec = batch_size * step / timer.elapsed()
-
-            # Save checkpoint every save_interval minutes
-            if(int(timer.elapsed()) >= 60 * save_interval):
-                timer.reset()
-
-                save_diffusion_model(
-                    save_path,
-                    diffusion_prior,
-                    trainer.optimizer,
-                    trainer.scaler,
-                    config,
-                    image_embed_dim)
-
-            # Log to wandb
-            tracker.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,
-                        image_reader,
-                        text_reader,
-                        train_set_size,
-                        NUM_TEST_EMBEDDINGS,
-                        device)
-                ### Evaluate model(validation run) ###
-                eval_model(diffusion_prior,
-                        device,
-                        image_reader,
-                        text_reader,
-                        eval_start,
-                        eval_start+NUM_TEST_EMBEDDINGS,
-                        NUM_TEST_EMBEDDINGS,
-                        dp_loss_type,
-                        phase="Validation")
-
-            trainer.update()
-
-    ### Test run ###
-    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")
-
-if __name__ == "__main__":
-    train()