How to implement Stable Diffusion

Import what we need

The first thing is to import the libraries and methods that we are going to use:

import torch
from torchvision.transforms import ToPILImage
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, EulerDiscreteScheduler
from tqdm.auto import tqdm
from PIL import Image

Later you will understand what each thing is for.

Instantiate models

We are going to instantiate the necessary models in order to have them available throughout the application.

We need the tokenizer, the transformer (text encoder), the noise predictor (U-Net), the scheduler and the variational autoencoder (VAE). All models are obtained from the Stable Diffusion 1.5 Hugging Face repository.

Each model is extracted from a specific folder (subfolder) and we make use of the .safetensors format when available. In addition, the parameterized models are moved to the graphics card using to('cuda'), to speed up the computations.

tokenizer = CLIPTokenizer.from_pretrained(
  'runwayml/stable-diffusion-v1-5',
  subfolder='tokenizer',
)

text_encoder = CLIPTextModel.from_pretrained(
  'runwayml/stable-diffusion-v1-5',
  subfolder='text_encoder',
  use_safetensors=True,
).to('cuda')

scheduler = EulerDiscreteScheduler.from_pretrained(
  'runwayml/stable-diffusion-v1-5',
  subfolder='scheduler',
)

unet = UNet2DConditionModel.from_pretrained(
  'runwayml/stable-diffusion-v1-5',
  subfolder='unet',
  use_safetensors=True,
).to('cuda')

vae = AutoencoderKL.from_pretrained(
  'runwayml/stable-diffusion-v1-5',
  subfolder='vae',
  use_safetensors=True,
).to('cuda')

Stable Diffusion 1.5 uses the scheduler PLMS (also called PNDM), but we are going to use Euler to show how easy it is to replace it (we have already done so).

Initialize parameters

Next, we define the parameters we need for image generation.

Let's define our prompt. We use a list in case we would like to generate several images at the same time using different prompts (['prompt1', 'prompt2', '...']) or several images of the same prompt using a random seed (['prompt'] * 4). For now we are not going to complicate things, we use a single prompt:

prompts = ['silly dog wearing a batman costume, funny, realistic, canon, award winning photography']

To make the code more readable, we store how many images we are going to generate at the same time:

batch_size = 1

For the sampling process, we specify that we are going to use 30 steps (steps or sampling steps). That is, the image will be denoised 30 times.

inference_steps = 30

To obtain a reproducible result we specify a seed instead of being random. This number will be used later to generate a tensor with noise from which we will clean the image until we obtain the result. If we start from the same noise, we will always get the same result.

seed = 1055747

And finally, we also save in some variables the value of CFG, as well as the size of the image we want to generate:

cfg_scale = 7
height = 512
width = 512

Conditioning

Let's start by generating the tensor that contains the information to guide the noise predictor towards the image we expect to obtain.

print(tokenizer('dog in batman costume'))

{'input_ids': [49406, 1929, 530, 7223, 7235, 49407], 'attention_mask': [1, 1, 1, 1, 1, 1]}

"<|startoftext|>": 49406,
"dog</w>": 1929,
"in</w>": 530,
"batman</w>": 7223,
"costume</w>": 7235,
"<|endoftext|>": 49407,

cond_input = tokenizer(
  prompts,
  max_length=tokenizer.model_max_length,   # Size we need (77)
  padding='max_length',                    # Apply padding if necessary
  truncation=True,                         # Apply truncation if necessary
  return_tensors='pt',
)

print(cond_input)

{'input_ids': tensor([[49406,  1929,   530,  7223,  7235, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0]])}

with torch.no_grad():
  cond_embeddings = text_encoder(cond_input.input_ids.to('cuda'))[0]

print(cond_embeddings)

tensor([[[-0.3884,  0.0229, -0.0522,  ..., -0.4899, -0.3066,  0.0675],
         [-1.8022,  0.5477,  1.0725,  ..., -1.5483, -0.5022, -0.2065],
         [-0.3402,  1.4715, -0.7753,  ..., -1.0974, -0.6557,  0.0747],
         ...,
         [-0.9392,  0.1777,  0.2575,  ...,  0.9130, -0.3660, -1.0388],
         [-0.9334,  0.1780,  0.2590,  ...,  0.9113, -0.3683, -1.0343],
         [-0.8973,  0.1848,  0.2609,  ...,  0.9189, -0.3297, -1.0798]]],
       device='cuda:0')

Unconditioning

Stable Diffusion must also be provided with an unconditioned embedding. It is done in the same way but the prompt is an empty string as many times as there are images we are generating at the same time.

uncond_input = tokenizer(
  [''] * batch_size,
  max_length=tokenizer.model_max_length,
  padding='max_length',
  truncation=True,
  return_tensors='pt',
)

with torch.no_grad():
  uncond_embeddings = text_encoder(uncond_input.input_ids.to('cuda'))[0]

We join these embeddings, both conditioned and unconditioned, into a single tensor:

text_embeddings = torch.cat([uncond_embeddings, cond_embeddings])

Generate a tensor with noise

To generate noise we instantiate a generator using torch.Generator and assign it the seed from which we will start:

generator = torch.Generator(device='cuda')
generator.manual_seed(seed)

Next, we use the function torch.randn to obtain a tensor with noise. The parameter it receives is a sequence of integers that defines the shape of the tensor:

latents = torch.randn(
  (batch_size, unet.config.in_channels, height // 8, width // 8),
  generator=generator,
  device='cuda',
)

We are passing (1, 4, 64, 64) as value. These numbers come from:

• 1: The value of batch_size. That is, how many images we generate at once.
• 4: The number of input channels that the noise predictor neural network (U-Net) has.
• 64/64: The size of the image in latent space (height / width). Our image is 512x512 in size but in latent space it occupies 8 times less. This divider is specified in the Stable Diffusion 1.5 architecture.

Remember that when using CUDA, we have specified it in both functions using the device='cuda' argument.

Now latents is a tensor with noise on which we can now work as if it were a canvas.

Clean up the tensor noise

Let's configure the scheduler to specify in how many steps we want to clean the tensor:

scheduler.set_timesteps(inference_steps)

We can check how it works internally by printing the following property:

print(scheduler.timesteps)

tensor([999.0000, 964.5517, 930.1035, 895.6552, 861.2069, 826.7586, 792.3104,
        757.8621, 723.4138, 688.9655, 654.5172, 620.0690, 585.6207, 551.1724,
        516.7241, 482.2758, 447.8276, 413.3793, 378.9310, 344.4828, 310.0345,
        275.5862, 241.1379, 206.6897, 172.2414, 137.7931, 103.3448,  68.8966,
         34.4483,   0.0000])

As there are 30 steps, 30 elements have been generated separated by the same distance (34.4483 units).

Some schedulers, such as DPM2 Karras or Euler, require that from the first step the tensor values are already multiplied by the standard deviation of the initial noise distribution. Schedulers that do not need it will simply multiply by 1. That is... we need to add the following operation:

latents = latents * scheduler.init_noise_sigma

We already have the tensor with noise, the conditioning and the scheduler. With this we can now generate the loop that will clean the tensor over 30 turns. We use the tqdm library to display a progress bar. I will explain each line directly in the code.

for t in tqdm(scheduler.timesteps):
  # Since we are using classifier-free guidance, we duplicate the tensor to avoid making two passes
  # One pass will be for the conditioned values and another for the unconditioned ones
  # With this it gains efficiency
  latent_model_input = torch.cat([latents] * 2)

  # This line ensures compatibility between various schedulers
  # Basically, the ones that need to scale the input based on the timestep
  latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)

  with torch.no_grad():
    # The U-Net is asked to make a prediction of the amount of noise in the tensor
    noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

  # We assign half of the estimated noise to conditioning and the other half to unconditioning
  noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)

  # Here the prediction is adjusted to guide towards a conditioned result
  # That is, more or less importance is given to the conditioning (to the prompt in this case)
  noise_pred = noise_pred_uncond + cfg_scale * (noise_pred_cond - noise_pred_uncond)

  # A new tensor is generated by subtracting the amount of noise we have previously calculated
  # This is the process that cleans the noise until all the scheduler steps have finished
  latents = scheduler.step(noise_pred, t, latents).prev_sample

After finishing this loop, we have our tensor free of noise and ready to be converted into a spectacular image. Or maybe in a churro, now we will find out.

latents = scheduler.step(noise_pred, t, latents, generator=generator).prev_sample

Converting the tensor into an image

There's little left. The scale factor (vae.config.scale_factor) of the VAE itself must be taken into account, a value set at 0.18215. Once the tensor is normalized we can decode it using vae.decode() to take it out of latent space and put it into image space.

latents = latents / vae.config.scale_factor
with torch.no_grad():
  images = vae.decode(latents).sample

We still have a tensor, but it is no longer inside the Matrix. This tensor has values ranging from -1 to 1, so we first normalize it to a range from 0 to 1.

We could perform calculations to convert these values into RGB values but let's allow torchvision to handle that. We use its ToPILImage transformation to convert the torch tensor into a pillow image (or multiple images, depending on the batch_size). The save() method from Pillow will take care of saving the images to disk.

images = (images / 2 + 0.5).clamp(0, 1)

to_pil = ToPILImage()

for i in range(1, batch_size + 1):
  image = to_pil(images[i])
  image.save(f'image_{i}.png')

Run python inference.py... and we can now see our gorgeous and perfect image!