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.
# Supress some unnecessary warnings when loading the CLIPTextModellogging.set_verbosity_error()
# Set devicetorch_device ="cuda"if torch.cuda.is_available()else"cpu"
Load Models
# Load the autoencoder model which will be used to decode the latents into image space. vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")# Load the tokenizer and text encoder to tokenize and encode the text. tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")# The UNet model for generating the latents.unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")# The noise schedulerscheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
# To the GPU we go!vae = vae.to(torch_device)text_encoder = text_encoder.to(torch_device)unet = unet.to(torch_device);
Diffusion Loop
We will dig deeper into how a pipeline is set up to run diffusion end to end.
prompt = ["A watercolor painting of an otter"]height =512# default height of Stable Diffusionwidth =512# default width of Stable Diffusionnum_inference_steps =30# Number of denoising stepsguidance_scale =7.5# Scale for classifier-free guidancegenerator = torch.manual_seed(32)# Seed generator to create the inital latent noisebatch_size =1
% Total % Received % Xferd Average Speed Time Time Time Current
Dload Upload Total Spent Left Speed
100 62145 100 62145 0 0 994k 0 --:--:-- --:--:-- --:--:-- 994k
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.
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.
noise = torch.randn_like(latent)# Random noisesampling_step =10# Equivalent to step 10 out of 15 in the schedule abovenoisy_latent = scheduler.add_noise(latent, noise, timesteps=torch.tensor([scheduler.timesteps[sampling_step]]))decode_latent_to_pil(noisy_latent)[0]
The scheduler's add_noise function simply adds noise scaled by sigma.
noisy_samples = original_samples + noise * 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:
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.
prompt = ["anime characte with a sword on his shoulder"]height =512# default height of Stable Diffusionwidth =512# default width of Stable Diffusionnum_inference_steps =30# Number of denoising stepsguidance_scale =8# Scale for classifier-free guidancegenerator = torch.manual_seed(3)# Seed generator to create the inital latent noisebatch_size =1# Prepare text (same as before)text_input =tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")with torch.no_grad(): text_embeddings =text_encoder(text_input.input_ids.to(torch_device))[0]max_length = text_input.input_ids.shape[-1]uncond_input =tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt")with torch.no_grad(): uncond_embeddings =text_encoder(uncond_input.input_ids.to(torch_device))[0]text_embeddings = torch.cat([uncond_embeddings, text_embeddings])# Prepare scheduler (setting the number of inference steps)scheduler.set_timesteps(num_inference_steps)# Prepare latents (noising appropriately for start_step)start_step =12start_sigma = scheduler.sigmas[start_step]noise = torch.randn_like(latent)noisy_latent = scheduler.add_noise(latent, noise, timesteps=torch.tensor([scheduler.timesteps[start_step]]))noisy_latent = noisy_latent.to(torch_device).float()# Loop and denoise# We start at later step to control how much noise we add to the original image.for i, t intqdm(enumerate(scheduler.timesteps)):if i >= start_step:# << This is the only modification to the loop we do# expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. latent_model_input = torch.cat([noisy_latent] *2) sigma = scheduler.sigmas[i] latent_model_input = scheduler.scale_model_input(latent_model_input, t)# predict the noise residualwith torch.no_grad(): noise_pred =unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]# perform guidance noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)# compute the previous noisy sample x_t -> x_t-1 noisy_latent = scheduler.step(noise_pred, t, noisy_latent).prev_sampledecode_latent_to_pil(noisy_latent)[0]
0it [00:00, ?it/s]
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.
prompt ='A picture of a puppy.'text_input =tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")print("Tokens:", text_input['input_ids'][0])# See the individual tokens# Any token after "puppy" is an end-of-text placeholder token.for t in text_input['input_ids'][0][:8]:print(t, tokenizer.decoder.get(int(t)))
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.
# Grab the output embeddingsoutput_embeddings =text_encoder(text_input.input_ids.to(torch_device))[0]print('Shape:', output_embeddings.shape)output_embeddings
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.
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.
Now if we put them together, this is what we get from the CLIP text encoder.
def get_output_embeds(input_embeddings):
# CLIP's text model uses causal mask, so we prepare it here:
batch_size, seq_len = input_embeddings.shape[:2]
causal_attention_mask = text_encoder.text_model._build_causal_attention_mask(batch_size,
seq_len,
dtype=input_embeddings.dtype)
# Getting the output embeddings involves calling the model with passing output_hidden_states=True
# so that it doesn't just return the pooled final predictions:
encoder_outputs = text_encoder.text_model.encoder(
inputs_embeds=input_embeddings,
attention_mask=None, # We aren't using an attention mask so that can be None
causal_attention_mask=causal_attention_mask.to(torch_device),
output_attentions=None,
output_hidden_states=True, # We want the output embs not the final output
return_dict=None,
)
# We're interested in the output hidden state only
output = encoder_outputs[0]
# There is a final layer norm we need to pass these through
output = text_encoder.text_model.final_layer_norm(output)
# And now they're ready!
return output
out_embs_test = get_output_embeds(input_embeddings) # Feed through the model with our new function
print(out_embs_test.shape) # Check the output shape
out_embs_test # Inspect the output
Here's the quick shortcut that will mix two sentences together.
text_input1 = tokenizer(["A giraffe with long neck"], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
text_input2 = tokenizer(["A cat with long neck"], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
text_embeddings1 = text_encoder(text_input1.input_ids.to(torch_device))[0]
text_embeddings2 = text_encoder(text_input2.input_ids.to(torch_device))[0]
mix_factor = 0.5
mixed_embeddings = text_embeddings1*mix_factor + text_embeddings2*(1-mix_factor)
# Generate!
generate_with_embeddings(mixed_embeddings)
0it [00:00, ?it/s]
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.
scheduler.set_timesteps(num_inference_steps)
t = scheduler.timesteps[0]
sigma = scheduler.sigmas[0]
print("time step", t, "and sigma", sigma)
# Let's take a random noise vector as the input to U-Net
noisy_latents = torch.randn((batch_size, unet.in_channels, height // 8, width // 8), generator=generator)
noisy_latents = noisy_latents.to(torch_device)
noisy_latents = noisy_latents * scheduler.init_noise_sigma
# Pick a prompt
text_input = tokenizer(['A macaw'],
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt")
with torch.no_grad():
text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
# Run it through U-Net
with torch.no_grad():
noise_pred = unet(noisy_latents, t, encoder_hidden_states=text_embeddings)["sample"]
noisy_latents.shape, noise_pred.shape
time step tensor(999., dtype=torch.float64) and sigma tensor(14.6146)
(torch.Size([1, 4, 64, 64]), torch.Size([1, 4, 64, 64]))
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.
prompt = 'Oil painting of an otter in a top hat'
prompt = 'K-pop boy band BTS is performing on stage'
height = 512
width = 512
num_inference_steps = 50
guidance_scale = 8
generator = torch.manual_seed(123)
batch_size = 1
# Make a folder to store results
!rm -rf steps/
!mkdir -p steps/
# Prepare text embeddings
text_input = tokenizer([prompt], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
max_length = text_input.input_ids.shape[-1]
# ? Why do we need unconditional inputs for classifier-free guidance?
uncond_input = tokenizer(
[""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
with torch.no_grad():
uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
# Prepare scheduler, use the noise to generate a latent
scheduler.set_timesteps(num_inference_steps)
latents = torch.randn((batch_size, unet.in_channels, height // 8, width // 8), generator=generator)
latents = latents.to(torch_device)
latents = latents * scheduler.init_noise_sigma
for i, t in tqdm(enumerate(scheduler.timesteps)):
# Expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
latent_model_input = torch.cat([latents] * 2)
sigma = scheduler.sigmas[i]
latent_model_input = scheduler.scale_model_input(latent_model_input, t)
# Predict the noise residual
with torch.no_grad():
noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]
# Perform guidance
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# Get the predicted x0:
# latents_x0 = latents - sigma * noise_pred # If we were to calculate it ourselves
latents_x0 = scheduler.step(noise_pred, t, latents).pred_original_sample # Using the scheduler (Diffusers 0.4 and above)
# Compute the previous noisy sample x_t -> x_t-1
latents = scheduler.step(noise_pred, t, latents).prev_sample
# To PIL Images
im_t0 = decode_latent_to_pil(latents_x0)[0]
im_next = decode_latent_to_pil(latents)[0]
# Combine the two images and save for later viewing
im = Image.new('RGB', (1024, 512))
im.paste(im_next, (0, 0))
im.paste(im_t0, (512, 0))
im.save(f'steps/{i:04}.jpeg')
0it [00:00, ?it/s]
# Make and show the progress video (change width to 1024 for full res)
!ffmpeg -v 1 -y -f image2 -framerate 12 -i steps/%04d.jpeg -c:v libx264 -preset slow -qp 18 -pix_fmt yuv420p out.mp4
from IPython.display import Video
Video("out.mp4")
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?
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.
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.
# ??scheduler.step
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.
def blue_loss(images):
# How far are the blue channel values to 0.9:
loss = torch.abs(images[:,2] - 0.9).mean() # [:,2] -> all images in batch, only the blue channel
return loss
We will update the latent in each time step based on this blue loss.
for i, t in tqdm(enumerate(scheduler.timesteps)):
# Expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
latent_model_input = torch.cat([latents] * 2)
sigma = scheduler.sigmas[i]
latent_model_input = scheduler.scale_model_input(latent_model_input, t)
# Predict the noise residual
with torch.no_grad():
noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]
# Perform classifier free guidance
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
#### Blue Guidance ###
if i%5 == 0:
# Requires grad on the latents
latents = latents.detach().requires_grad_()
# Get the predicted x0:
# latents_x0 = latents - sigma * noise_pred
latents_x0 = scheduler.step(noise_pred, t, latents).pred_original_sample
# Decode to image space
denoised_images = vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5 # range (0, 1)
# Calculate loss
loss = blue_loss(denoised_images) * blue_loss_scale
# Occasionally print it out
if i%10==0:
print(i, 'loss:', loss.item())
# Get gradient
cond_grad = torch.autograd.grad(loss, latents)[0]
# Modify the latents based on this gradient
latents = latents.detach() - cond_grad * sigma**2
# Now step with scheduler
latents = scheduler.step(noise_pred, t, latents).prev_sample
decode_latent_to_pil(latents)[0]
