readme: ---
  license: mit
  ---
  # A Text-Conditioned Diffusion-Prior
  
  ## Training Details
  
  [Updated Reports Coming]
  
  ## Source Code
  Models are diffusion prior trainers from https://github.com/lucidrains/DALLE2-pytorch
  
  ## Community: LAION
  Join Us!: https://discord.gg/uPMftTmrvS
  
  ---
  
  ## Intro
  
  A properly trained prior will allow you to translate between two embedding spaces. If you know *a priori* that two embeddings are connected some way—then ability the translate between them could extremely helpful.
  
  ### Motivation
  
  Before we dive into the model, let’s look at a quick example of where the model may be helpful.
  
  For demonstration purposes we will imagine that we wish to generate images from text using CLIP and a Decoder.
  
  > [CLIP](https://openai.com/blog/clip/) is a contrastive model that learns to maximize the cosine similarity between a given image and caption, however, there is no guarantee that these embeddings are in the same space. While the embeddings generated are ***close*** the image and text embeddings occupy two disjoint sets.
  
  ```python
  # Load Models
  clip_model = clip.load("ViT-L/14")
  decoder = Decoder(checkpoint="best.pth") # A decoder trained on CLIP Image embeddings
  
  # Retrieve prompt from user and encode with CLIP
  prompt = "A corgi wearing sunglasses"
  tokenized_text = tokenize(prompt)
  text_embedding = clip_model.encode_text(tokenized_text)
  
  # Now, pass the text embedding to the decoder
  predicted_image = decoder.sample(text_embedding)
  ```
  
  > **Question**: *Can you spot the issue here?*
  >
  > **Answer**: *We’re trying to generate an image from a text embedding!*
  
  Unfortunately, we run into the issue previously mentioned--the image embeddings and the text embeddings are not interchangeable! Now let's look at a better solution
  
  ```python
  # Load Models
  prior= Prior(checkpoint="prior.pth") # A decoder trained to go from: text-> clip text emb -> clip img emb
  decoder = Decoder(checkpoint="decoder.pth") # A decoder trained on CLIP Image embeddings
  
  # Retrieve prompt from user and encode with a prior
  prompt = "A corgi wearing sunglasses"
  tokenized_text = tokenize(prompt)
  text_embedding = prior.sample(tokenized_text) # <-- now we get an embedding in the same space as images!
  
  # Now, pass the predicted image embedding to the decoder
  predicted_image = decoder.sample(text_embedding)
  ```
  
  With the prior we are able to successfully generate embeddings *within* CLIP's image space! For this reason, the decoder will perform much better as it receives input that is much closer to its training data.
  
  > **You may be asking yourself the following question:**
  >
  > *"Why don't you just train the decoder on clip text embeddings instead of image embeddings?"*
  >
  > OpenAI covers this topic in their [DALLE-2 paper](https://arxiv.org/abs/2204.06125). The TL;DR is *"it doesn't work as well as decoders trained on image embeddings"*...also...its just an example :smile:
  
  ## Usage
  
  To utilize a pre-trained prior, it’s quite simple.
  
  ### Loading Checkpoints
  ```python
  import torch
  from dalle2_pytorch import DiffusionPrior, DiffusionPriorNetwork, OpenAIClipAdapter
  from dalle2_pytorch.trainer import DiffusionPriorTrainer
  
  def load_diffusion_model(dprior_path):
  
      prior_network = DiffusionPriorNetwork(
          dim=768,
          depth=24,
          dim_head=64,
          heads=32,
          normformer=True,
          attn_dropout=5e-2,
          ff_dropout=5e-2,
          num_time_embeds=1,
          num_image_embeds=1,
          num_text_embeds=1,
          num_timesteps=1000,
          ff_mult=4
      )
  
      diffusion_prior = DiffusionPrior(
          net=prior_network,
          clip=OpenAIClipAdapter("ViT-L/14"),
          image_embed_dim=768,
          timesteps=1000,
          cond_drop_prob=0.1,
          loss_type="l2",
          condition_on_text_encodings=True,
  
      )
  
      trainer = DiffusionPriorTrainer(
          diffusion_prior=diffusion_prior,
          lr=1.1e-4,
          wd=6.02e-2,
          max_grad_norm=0.5,
          amp=False,
          group_wd_params=True,
          use_ema=True,
          device=device,
          accelerator=None,
      )
  
      trainer.load(dprior_path)
  
      return trainer
  ```
  
   Here we instantiate a model matches the configuration it was trained with, and then load the weights (*just like any other PyTorch model!*)
  
  ### Sampling
  Once we have a pre-trained model, generating embeddings is quite simple!
  ```python
  # tokenize the text
  tokenized_text = clip.tokenize("<your amazing prompt>")
  # predict an embedding
  predicted_embedding = prior.sample(tokenized_text, n_samples_per_batch=2, cond_scale=1.0)
  ```
  
  The resulting tensor returned from `.sample()` is of the same shape as your training data along the non-batch dimension(s). For example, a prior trained on `ViT-L/14` embeddings will predict an embedding of shape (1, 768).
  
  > For CLIP priors, this is quite handy as it means that you can use prior.sample(tokenizer_text) as a drop in replacement for clip.encode_text().
  
  **Some things to note:**
  * It is possible to specify the number of embeddings to sample from