Stable Diffusion Deep Dive
Examples are provided by fastai's stable diffusion GitHub repository.
Setup
We will switch gear and use PyTorch instead of TensorFlow for the following sections. We also need a huggingface account for loading the pre-trained models.
# !pip install -q --upgrade transformers diffusers ftfyfrom base64 import b64encode
from pathlib import Path
import numpy
import torch
from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
from huggingface_hub import notebook_login# For video display:
from IPython.display import HTML
from matplotlib import pyplot as plt
from PIL import Image
from torch import autocast
from torchvision import transforms as tfms
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer, loggingLoad Models
Diffusion Loop
We will dig deeper into how a pipeline is set up to run diffusion end to end.
Embed text input
Scheduler
Latent vectors are going to be fed into U-Net, not the actual image itself.
Start the loop.
Variational Autoencoder
The VAE encodes images into latent vectors before sending them to U-Net for generation step.
Download a demo image
Encode the image into latent space, i.e. going from (512, 512, 3) to (64, 64, 4).
We can inspect what do the latents look like, dimension by dimension.
The latents are acting as lossy compressed images.
The difference is very small. Look at the eye and see the small difference. This is effectively a 48x compression going from input to latent. Thus instead of doing diffusion on image space, we can make it work on latent space because latent space does an excellent job at capturing image information. With 48x reduction, diffusion can be run on much cheaper hardware. This is one of the key innovations of stable diffusion.
Scheduler
The training objective for a diffusion model is to "predict" noise for a given image. We then subtract the noise from the image to obtain a final image output. When the model is fully trained, we can feed it pure noise and model will output all the noise it sees. When we subtract the noise from a pure noise image, a recognizable image will emerge.
During training time, we need to vary the amount of noise added for every training step. If too much noise, the model doesn't have much information to learn from. If too little noise, the model loses generative ability. We need to follow a schedule to add noise to each training step according to some distribution.
During sampling time or inference time, we need to vary the amount of "denoise" for every inference step. We need to make use of the same schedule that we used during training.
This is where the linear multi-step discerete scheduler comes in.
During sampling, we will start at a high noise level (in fact, our input will be pure noise) and gradually 'denoise' down to an image.
Sigma is the amount of noise added to the latent representation. Let's add some noise to an encoded image and see what the decoding will look like.
The scheduler's add_noise function simply adds noise scaled by sigma.
Other diffusion models may be trained with different noising and scheduling approaches, some of which keep the variance fairly constant across noise levels ('variance preserving') with different scaling and mixing tricks instead of having noisy latents with higher and higher variance as more noise is added ('variance exploding').
If we want to start from random noise instead of a noised image, we need to scale it by the largest sigma value used during training, ~14 in this case. And before these noisy latents are fed to the model they are scaled again in the so-called pre-conditioning step:
now handled by scheduler's scale_model_input API.
Again, this scaling/pre-conditioning differs between papers and implementations, so keep an eye out for this if you work with a different type of diffusion model.
Image2Image Architecture
Instead of using pure noise as starting point, we use an actual image and add noise to it. This will produce a new image that is based on the input image with variations.
Text Embedding via CLIP Encoder
The tokenizer will embed the text prompt as a vector. It will first turn the text into tokens and then tokens become embedding vectors.
Now we look at the actual embeddings from the text encoder. The text encoder we use is called CLIPTextModel which stands for contrastive language-image pre-training text model. This model is trained on pairing text captions with images.
There are two steps to embedding, token embedding and position embedding.
Token Embeddings
We have a vocab size of 49408 and embedding dimension 768.
The embedding is simply an index look-up.
We can do it to the whole sentence.
Positional Embedding
Word order matters. We can embed each word position as a vector as well. In fact the text model expects us to provide both token and position embeddings to the encoder. Our sequence max length is 77 so we should expect position embedding to return 77 at axis=0.
Combining Token and Position Embeddings
We can sum the two vectors for this particular text encoding model to create a final embedding. In other models, the two vectors may be concatenated to form the final embedding vector.
The operation is identical to the following.
Now if we put them together, this is what we get from the CLIP text encoder.
It is identical to the following operation.
Modify Embeddings
Let's see what happens if we modify one of the tokens.
Generate an image with this new embedding.
Mixed Text Encoding
We can mix two different prompt together and assign a weight ratio to them. Let's make a panda cat. First we need to know the token for cat and panda.
Here's the quick shortcut that will mix two sentences together.
UNet
We will look into how the de-noising actually works. U-Net takes the latent vector and text embeddings to produce a new latent vector which then decode into image.
Given a set of noisy latents, the model predicts the noise component. We subtract the noise from the noisy latents to generate a denoised latent that can be interpreted as an image, like this.
Let's visualize the progression of denoising.
Your browser does not support the video element.
The left image is the result of decoding somewhat-noisy latent for every time step. The right image is the result of decoding a fully noise-subtracted latent for every time step. Note that scheduler does not subtract all noises from latent when it produces prev_sample. More on this later.
Classifier Free Guidance
We first need to ask what is a classifier guidance?
Excerpt from Classifier-Free Diffusion Guidance:
Dhariwal & Nichol (2021) proposed a classifier guidance, a technique to boost the sample quality of a diffusion model using an extra trained classifier. Prior to classifier guidance, it was not known how to generate low temperature samples from a diffusion model... Classifier guidance instead mixes a diffusion model's score estimate with the input gradient of the log probability of a classifier.
Low temperature samples refer to denoised samples (in physics, low temperature = high entropy, high temperature = low entropy). The classifier provides "guidance" to the de-noising procedure.
The classifier-free approach is an alternative method of modifying conditional embedding to have the same effect as classifier guidance. That is to use a pair of conditional and unconditional text embedding as a pair of positive/negative sample. The unconditional text embedding is just a sentence of null/placeholder tokens.
The text embedding is fed into U-Net as a conditional
We can break the prediction into two parts, one is based on conditional, another one is based on unconditional. We can break them into two parts.
The final noise latent prediction is then a linear combination of the two, weighted by guidance scale.
During training time, the model learns to distinguish conditional vs unconditional. It dcreases the unconditional likelihood of the sample while increasing the conditional likelihood. It is effectively learning to be a classifier.
Sampling
The model tries to predict noise in an image. For low noise values, it does a pretty good job. For higher noise levels, it struggles. Instead of producing a sharp image, the results tend to look like a blurry mess. Samplers use model predictions to move a small amount towards the model prediction by remving some noises but not all, then get another prediction based on this marginally-less-rubbish input. Different scheduler/sampler does this in different ways.
Mess with Guidance
Taking the classifier guidance idea, we can introduce a gradient to guide the image generation process. For example, we want a blue image. We will define a loss function with respect to noisy latents. We take the gradient and use the gradient to modify latent.
We will update the latent in each time step based on this blue loss.
Last updated