Cómo implementar Stable Diffusion

Importar lo necesario

Lo primero es importar las librerías y métodos que vamos a utilizar:

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

Más adelante entenderás para qué es cada cosa.

Instanciar modelos

Vamos a instanciar los modelos necesarios para así tenerlos disponibles en toda la aplicación.

Necesitamos el tokenizador, el transformer (text encoder), el predictor de ruido (U-Net), el scheduler y el variational autoencoder (VAE). Todos los modelos los obtenemos del repositorio en Hugging Face de Stable Diffusion 1.5.

Cada modelo se extrae de una carpeta específica (subfolder) y hacemos uso del formato .safetensors cuando esté disponible. Además, los modelos parametrizados los movemos a la tarjeta gráfica mediante to('cuda'), para acelerar los cálculos.

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 utiliza el scheduler PLMS (también llamado PNDM), pero nosotros vamos a utilizar Euler para mostrar lo fácil que es sustituirlo (ya lo hemos hecho).

Inicializar parámetros

A continuación, definimos los parámetros que necesitamos para la generación de imágenes.

Vamos a definir nuestro prompt. Utilizamos una lista por si quisiéramos generar varias imágenes al mismo tiempo utilizando diferentes prompts (['prompt1', 'prompt2', '...']) o varias imágenes del mismo prompt utilizando una semilla aleatoria (['prompt'] * 4). De momento no nos vamos a complicar, utilizamos un único prompt:

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

Para hacer el código más legible, almacenamos cuántas imágenes vamos a generar al mismo tiempo:

batch_size = 1

Para el proceso de muestreo, especificamos que vamos a utilizar 30 pasos (steps o sampling steps). Es decir, se eliminará ruido de la imagen 30 veces.

inference_steps = 30

Para obtener un resultado reproducible especificamos una semilla en vez de ser aleatoria. Este número se utilizará más adelante para generar un tensor con ruido desde el que iremos limpiando la imagen hasta obtener el resultado. Si partimos del mismo ruido, siempre obtendremos el mismo resultado.

seed = 1055747

Y por último, guardamos también en unas variables el valor de CFG, así como el tamaño de la imagen que queremos generar:

cfg_scale = 7
height = 512
width = 512

Condicionamiento

Empecemos generando el tensor que contiene la información para guiar al predictor de ruido hacia la imagen que esperamos obtener.

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,   # Tamaño que necesitamos (77)
  padding='max_length',                    # Aplicar padding si fuera necesario
  truncation=True,                         # Aplicar truncation si fuera necesario
  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')

Incondicionamiento

A Stable Diffusion tambien hay que proveerle de un embedding no condicionado. Se hace de la misma manera pero el prompt es una cadena vacía tantas veces como imágenes estemos generando a la vez.

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]

Unimos estos embeddings, tanto condicionados como no condicionados, en un único tensor:

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

Generar un tensor con ruido

Para generar ruido instanciamos un generador mediante torch.Generator y le asignamos la semilla desde la que comenzaremos:

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

Después, utilizamos la función torch.randn para obtener un tensor con ruido. El parámetro que recibe es una secuencia de enteros que define la forma del tensor:

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

Le estamos pasando como valor (1, 4, 64, 64). Estos números provienen de:

• 1: El valor de batch_size. Es decir, cuántas imágenes generamos a la vez.
• 4: La cantidad de canales de entrada que tiene la red neuronal del predictor de ruido (U-Net).
• 64/64: El tamaño de la imagen en el espacio latente (height / width). Nuestra imagen es de tamaño 512x512 pero en el espacio latente ocupa 8 veces menos. Este divisor viene especificado en la arquitectura de Stable Diffusion 1.5.

Recuerda que al utilizar CUDA, lo hemos especificado en ambas funciones mediante el argumento device='cuda'.

Ahora latents es un tensor con ruido sobre el que ya podemos trabajar como si de un lienzo se tratase.

Limpiar el ruido del tensor

Vamos a configurar el scheduler para indicarle en cuántos pasos queremos limpiar el tensor:

scheduler.set_timesteps(inference_steps)

Podemos comprobar cómo funciona internamente imprimiendo la siguiente propiedad:

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])

Como son 30 pasos, se han generado 30 elementos separados por la misma distancia (34.4483 unidades).

Algunos schedulers como DPM2 Karras o Euler, necesitan que desde el primer paso los valores del tensor ya estén multiplicados por la desviación estándar de la distribución de ruido inicial. Los schedulers que no lo necesitan simplemente multiplicarán por 1. El caso... que hay que añadir la siguiente operación:

latents = latents * scheduler.init_noise_sigma

Ya tenemos el tensor con ruido, los condicionamientos y el scheduler. Con esto ya podemos generar el bucle que limpiará el tensor a lo largo de 30 vueltas. Utilizamos la librería tqdm para mostrar una barra de progreso. Explicaré cada línea directamente en el código.

for t in tqdm(scheduler.timesteps):
  # Como estamos utilizando classifier-free guidance, duplicamos el tensor para evitar hacer dos pasadas
  # Una pasada será por los valores condicionados y otra por los incondicionados
  # Con esto se gana eficiencia
  latent_model_input = torch.cat([latents] * 2)

  # Esta línea asegura compatibilidad entre varios schedulers
  # Básicamente, los que necesitan escalar la entrada según el timestep
  latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)

  with torch.no_grad():
    # Se le pide a la U-Net que haga una predicción de la cantidad de ruido que hay en el tensor
    noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

  # Asignamos la mitad del ruido estimado al condicionamiento y la otra mitad al incondicionamiento
  noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)

  # Aquí se ajusta la predicción para guiar hacia un resultado condicionado
  # Es decir, se le da más o menos importancia al condicionamiento (al prompt en este caso)
  noise_pred = noise_pred_uncond + cfg_scale * (noise_pred_cond - noise_pred_uncond)

  # Se genera un nuevo tensor restando la cantidad de ruido que hemos calculado previamente
  # Este es el proceso que limpia el ruido hasta que todos los pasos del scheduler han finalizado
  latents = scheduler.step(noise_pred, t, latents).prev_sample

Tras finalizar este bucle, ya tenemos nuestro tensor libre de ruido y listo para ser convertido en una espectacular imagen. O quizás en un churro, ahora saldremos de dudas.

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

Convertir el tensor en una imagen

Ya queda poco. Hay que tener en cuenta el factor de escala (vae.config.scale_factor) del propio VAE, un valor fijado en 0.18215. Una vez normalizado el tensor ya podemos decodificarlo mediante vae.decode() para sacarlo del espacio latente e introducirlo en el espacio de imagen.

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

Seguimos teniendo un tensor, pero ya no está dentro de la Matrix. Este tensor tiene valores que van desde -1 a 1, así que primero lo normalizamos a un rango desde 0 hasta 1.

Podríamos hacer cálculos para convertir estos valores en valores RGB pero dejemos que torchvision se encargue de ello. Utilizamos su transformación ToPILImage para convertir el tensor de torch en una imagen de pillow (o varias, dependiendo del batch_size). El método save() de Pillow se encargará de guardar las imágenes en el disco.

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')

Ejecuta python inference.py... ¡y ya podemos visualizar nuesta magnífica y perfectísima imagen!