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.
defget_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 outputout_embs_test =get_output_embeds(input_embeddings)# Feed through the model with our new functionprint(out_embs_test.shape)# Check the output shapeout_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.5mixed_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-Netnoisy_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 prompttext_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-Netwith 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 =512width =512num_inference_steps =50guidance_scale =8generator = torch.manual_seed(123)batch_size =1# Make a folder to store results!rm -rf steps/!mkdir -p steps/
# Prepare text embeddingstext_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 latentscheduler.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 intqdm(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 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)# 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 VideoVideo("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.
defblue_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 channelreturn loss
We will update the latent in each time step based on this blue loss.
for i, t intqdm(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 residualwith 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 outif 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_sampledecode_latent_to_pil(latents)[0]
