Cómo funciona Stable Diffusion
Comprende de manera sencilla cómo Stable Diffusion transforma unas pocas palabras en una espectacular imagen
Introducción
Para utilizar Stable Diffusion no hace falta entender cómo funciona, pero ayuda. Muchos parámetros configurables como Seed, Sampler, Steps, CFG Scale o Denoising strength son esenciales para generar imágenes de calidad.
En este artículo vamos a ver, a grandes rasgos, sin una sola línea de código y casi para todos los públicos, como funciona Stable Diffusion internamente en su versión 1.x y cómo es capaz de generar imágenes tan creativas con una simple frase de texto.
He publicado una segunda parte de este artículo. En este nuevo artículo implementaremos Stable Diffusion 1.5 utilizando código en Python. Ya que aquí veremos la teoría detrás de la generación de imágenes, en la segunda parte lo pondremos en práctica y veremos cómo están interconectadas las piezas entre sí. No es porque lo diga yo, pero es un complemento ideal a este artículo.
Cómo funciona un modelo de difusión
Dentro del aprendizaje automático hay varios tipos de modelos generativos (modelos capaces de generar datos, en este caso imágenes). Una categoría son los modelos de difusión y se llaman así porque imitan el comportamiento de la difusión molecular. Un fenómeno que todos hemos experimentado cuando utilizamos ambientador o echamos un hielo en un vaso de agua. Las particulas tienden a moverse y esparcirse.
Difusión hacia delante
En el entrenamiento de Stable Diffusion se parte de cientos de miles de fotografías sacadas de Internet y se comienza añadiendo ruido gaussiano a las imágenes a lo largo de una serie de pasos hasta que las imágenes pierden todo el significado. Es decir, se difunde ruido en las imágenes igual que un ambientador difunde sus partículas en una habitación. A este proceso se le conoce como forward diffusion (ya que va hacia delante).
Difusión inversa
Cuando pedimos a Stable Diffusion que genere una imagen, sucede la magia conocida como reverse diffusion, un proceso que revierte el visto anteriormente.
También se utiliza el término inferencia (inference) cuando se pone en práctica lo que un modelo ha aprendido durante el entrenamiento. En el caso de Stable Diffusion se puede utilizar este término para el proceso de difusión reversa.
El primer paso es generar una imagen de 512x512 pixeles llena de ruido aleatorio, una imagen sin ningún significado. Para generar esta imagen llena de ruido también podemos modificar un parámetro conocido como semilla (seed), cuyo valor por defecto es -1 (aleatorio). Por ejemplo, si utilizamos la semilla número 4376845645 siempre se generará el mismo ruido, de ahí que se pueda replicar el resultado de una imagen si conocemos su semilla.
Después, para convertir una imagen llena de ruido en un perro disfrazado de Batman y revertir el proceso, hay que saber cuánto ruido hay en la imagen. Esto se hace entrenando a una red neuronal convolucional para que sea capaz de predecir el nivel de ruido dentro de una imagen. Este modelo se llama U-Net (o denoising U-Net), aunque en Stable Diffusion lo conocemos como predictor de ruido (noise predictor).
El entrenamiento de este modelo consiste en utilizar las imágenes del proceso de difusión hacia delante, indicando además cuánto ruido hemos añadido para que sea capaz de aprenderlo. Si te digo que la primera imagen tiene 0% de ruido, la del medio un 50% y la última un 100%... seguro que serías capaz de distinguir el nivel de ruido existente en una nueva fotografía porque has entendido el patrón. Quizás no a la perfección pero eh, ¡tú no eres una red neuronal!
Muestreo
Con esta imagen inicial llena de ruido y la manera de calcular cuánto ruido queda aún en la imagen comienza el proceso de muestro (sampling).
El predictor de ruido estima cuánto ruido hay en la imagen. Tras esto, el algoritmo llamado sampler genera una imagen con esa cantidad de rudio y se resta de la imagen original. Este proceso se repite la cantidad de veces especificada por los pasos (steps o sampling steps).
Algunos ejemplos de algoritmos de sampling son Euler, Euler Ancestral, DDIM, DPM, DPM2 o DPM++ 2M Karras. La diferencia está en que unos son más rápidos, otros más creativos, otros perfeccionan mejor los pequeños detalles, etc. Más información sobre estos algoritmos, sus usos y diferencias, en el artículo Guía completa de samplers en Stable Diffusion.
Después de repetir este proceso las veces necesarias obtendremos nuestra imagen final.
Otra pieza importante en el muestreo es el noise scheduler ya que tiene la función de gestionar cuánto ruido hay que eliminar en cada paso. Si la reducción de ruido fuese lineal, nuestra imagen cambiaría la misma cantidad en cada paso, produciendo cambios bruscos. Un noise scheduler con pendiente negativa puede eliminar grandes cantidades de ruido al inicio y así avanzar más rápido, para después pasar a eliminar menos cantidad de ruido y así afinar los pequeños detalles de la imagen. Este algoritmo se puede configurar y elegir entre múltiples opciones.
Pero... ¿cómo sabe el modelo que detrás del ruido hay un perro? La respuesta es que no lo sabe. De momento hemos conseguido generar una imagen no condicionada. El resultado puede ser un coche o cualquier cosa porque nuestro prompt no está condicionando el resultado.
Condicionamientos
Lo interesante, por supuesto, no es generar imágenes aleatorias sin sentido, si no condicionar el resultado en cada paso.
Si recordamos, hemos entrenado al predictor de ruido dándole dos cosas: la imagen y el nivel de ruido. Por lo tanto podemos dirigir al predictor de ruido hacia el resultado que queremos. Es como preguntar "¿cuál es el nivel de ruido en esta fotografía de un perro disfrazado de Batman?", en vez de simplemente preguntar "¿cuál es el nivel de ruido?". De la primera manera condicionamos el resultado.
¡No pienses en un elefante rosa!
Utilizar texto (prompts) no es la única manera de condicionar al predictor de ruido. Se pueden utilizar varios métodos e incluso múltiples al mismo tiempo. Veamos algunos ejemplos comenzando por el más importante.
Text-to-Image
Convertir texto en un condicionamiento es la función más básica de Stable Diffusion. Se le conoce como Text-to-Image y consta de varias piezas que vamos a ver a continuación.
Tokenizer
Como los ordenadores no entienden de letras, la primera tarea es utilizar un tokenizador (tokenizer) para convertir cada palabra en un número llamado símbolo (token).
En el caso de Stable Diffusion la tokenización se realiza gracias al modelo CLIP (Contrastive Language-Image Pre-Training). Este modelo creado por OpenAI convierte imágenes en descripciones detalladas, aunque para esta tarea solo se utiliza el tokenizador que incorpora.
El tokenizador no siempre convierte cada palabra en un token. Puede haber palabras compuestas y el uso del espacio también afecta al número de tokens.
Quizás hayas visto en algún lado que el máximo de palabras que puede tener un prompt es de 77. Esta limitación existe porque los tokens se guardan en un vector que tiene un tamaño de 77 tokens (1x77). Por lo tanto, esta limitación en realidad se refiere a los tokens, no a las palabras. Y aunque si bien es cierto que después de 77 tokens el resto se eliminan, actualmente se utilizan otras técnicas de concatenación y retroalimentación para condicionar al predictor de ruido con tantos bloques de 77 tokens como sea necesario hasta utilizar todos.
Este paso es el único exclusivo de Text-to-Image, a partir de ahora todos los condicionamientos que veremos después utilizan las siguientes capas.
Embedding
Ahora vamos a realizar un ejercicio de imaginación. Tenemos 10 fotografías de personas y tenemos que clasificarlas en un gráfico de dos coordenadas según dos parámetros: edad y cantidad de pelo. Si hacemos bien la tarea obtendremos algo como lo siguiente:
El gráfico contiene dos dimensiones (X e Y) para representar los grupos.
Pues bien, esto mismo es un embedding. Una manera de clasificar términos y conceptos almacenando los datos en vectores. En el caso de Stable Diffusion, se utiliza una versión de CLIP llamada ViT-L que contiene 768 dimensiones. No se puede saber cuales son esas dimensiones porque dependen de los datos de entrenamiento, pero podríamos imaginar que una dimensión sirve para representar el color, otra el tamaño de los objetos, otra para la textura, otra para la expresión facial, otra para la luminosidad, otra para la distancia espacial entre objectos... y así hasta 768.
Se utilizan embeddings porque una vez que se han categorizado los términos en dimensiones, el modelo es capaz de calcular la distancia entre ellos.
Mediante esta técnica, la dimensión que agrupa el concepto de ropa estará más cerca de la dimensión que almacena valores como hombre y mujer, que de la dimensión que agrupa las partes de una estancia como ventana o puerta. Este ejemplo es bastante simple, en realidad todo ocurre de manera mucho más abstracta.
Cabe destacar que cada token contendrá 768 dimensiones. Es decir, si utilizamos la palabra coche en nuestro prompt, ese token se convertirá en un vector de 768 dimensiones. Una vez se realiza esto con todos los tokens tendremos un embedding de tamaño 1x77x768.
De esta manera se pueden medir las distancias de todas las dimensiones de todos los tokens.
Transformer
Este es el último paso del condicionamiento. En esta pieza se procesan los embeddings mediante un modelo transformer de CLIP.
Esta arquitectura de red neuronal se compone de múltiples capas y se encarga de condicionar al predictor de ruido para guiarle en cada paso hacia una imagen que represente la información de los embeddings.
En Text-to-Image los embeddings se llaman text embeddings por razones obvias, pero el modelo de transformer puede obtener embeddings de otros condicionamientos como veremos un poco más abajo.
Aquí viene lo interesante del transformer. Se conoce como self-attention cuando todos los embeddings tienen el mismo peso y afectan por igual al resultado. Además de esta arquitectura, dentro de estas capas Stable Diffusion también utiliza cross-attention. Esta técnica consiste en calcular las relaciones de dependencia entre los embeddings y producir un tensor con los resultados. Aquí es donde se tienen en cuenta las distancias entre las dimensiones.
Por ejemplo si utilizamos el prompt una pared roja con una puerta de madera, una arquitectura solo de self-attention podría generar una pared de madera con una puerta roja. Es totalmente válido pero posiblemente incorrecto.
Gracias al cross-attention, el modelo calcula la distancia entre madera/puerta y entre madera/pared y sabe que la primera opción es más probable/cercana. Esto lo hacemos todos los días, si te digo que pienses en un coche rojo seguramente estás pensando en un deportivo, un Ferrari tal vez, no en un Toyota Prius. Esto se debe a que en nuestro cerebro las palabras coche y rojo están más cerca de Ferrari que de Prius. A esta técnica se la conoce como atención (attention) porque imita el proceso cognitivo que experimentamos cuando pestamos atención a ciertos detalles.
Utilizando esta información, se guía al predictor de ruido para que vaya acercándose en cada paso a una imagen que contenga la estructura más probable dado un condicionamiento.
Etiquetas de clase y el valor CFG Scale
Este otro condicionamiento es importante porque ayuda al transformer utilizando etiquetas de clasificación.
Cuando se entrena un modelo de Stable Diffusion, se utiliza un clasificador externo (CG - classifier guidance). Este es un modelo que se encarga de asignar etiquetas (labels) a las imágenes de entrenamiento. Así pues las fotografías de perros tendrán asociadas las etiquetas de animal y perros, mientras que las fotografías de gatos tendrán las etiquetas de animal y gatos (entre muchas otras).
Cuando utilizamos un valor alto en la escala de CG, estamos pidiendo al transformer que aparte las etiquetas que no son similares y se centre en las que sí lo son. En cambio, con un valor bajo, es como si se acercasen las etiquetas y al pedir una fotografía de un gato, podríamos obtener la de cualquier otro animal (o cualquier otra etiqueta cercana).
Para evitar tener que utilizar un modelo clasificador externo y así reducir la dificultad durante el entrenamiento, el proceso se mejoró con una técnica llamada classifier-free guidance (CFG). Esto quiere decir que el guiamiento ya no necesita un modelo clasificador que se encargue de añadir etiquetas a las imágenes. Durante el entrenamiento de Stable Diffusion se utiliza la descripción de la imagen para ajustar las clases automáticamente.
Ajustando el parámetro CFG Scale podemos demarcar más o menos el territorio de las etiquetas que queremos utilizar para generar la imagen. Por eso se suele decir que este valor añade o quita libertad de generación. A mayor valor, más fidelidad porque se seleccionarán las etiquetas cuidadosamente. A menor valor, más libertad para que Stable Diffusion pueda generar algo, que aunque no sea tan fiel a nuestro prompt, si pueda conseguir mayor calidad al no tener tantas limitaciones.
En mi experiencia el valor óptimo está entre 5.5 y 7.5, aunque a veces toca subirlo o bajarlo (sobre todo aumentarlo cuando no se respeta el prompt).
Image-to-Image
Si se utiliza una imagen para condicionar el resultado estaremos aplicando Image-to-Image.
Por ejemplo, vamos a dibujar una casa de la manera más chapuceramente posible y le vamos a pasar a Stable Diffusion tanto esta imagen como el prompt: una fotografía realista de una casa de madera con una puerta roja en medio del bosque. Estos dos condicionamientos se aplicarán y obtendremos el siguiente resultado.
¿Cómo ha sido posible? Bien, ¿recuerdas que Stable Diffusion primero crea una imagen llena de ruido sobre la que trabajar? Pues en vez de esto, se añade una cierta cantidad de ruido a nuestra imagen de partida y ya tenemos la base.
Gracias al parámetro Denoising strenght podemos controlar cuanta cantidad de ruido queremos añadir. Si utilizamos el valor de 0 no se añadirá ruido y por lo tanto obtendremos una imagen idéntica a nuestro boceto. En cambio, con un valor de 1 convertiremos nuestro condicionamiento en puro ruido y el boceto ya no afectará al resultado (sería como volver a Text-to-Image). En alguna parte entre medias está el valor idóneo, depende de la libertad que queramos dar a Stable Diffusion para que respete nuestro input o no.
Inpainting
La técnica de Inpainting en la que podemos regenerar una zona de la imagen es un image-to-image aplicado a una zona específica.
Outpainting
De igual modo, la técnica de Outpainting sobre la que extendemos una imagen creando nuevas zonas que antes no existían, consiste en condicionar al predictor de ruido con varias capas. Generalmente se utiliza ruido para la nueva zona (ya que está vacía), el prompt y también parte de la imagen original (o toda).
Depth-to-Image
Para utilizar la información de tridimensionalidad de una imagen se utiliza un mapa de profundidad. Esta técnica se conoce como Depth-to-Image.
Un modelo como MiDaS procesa la imagen para extraer un mapa de profundidad. Después se añade ruido dependiendo del valor de Denoising strenght y esta información se convierte a un embedding. Junto con el resto de embeddings se condiciona el resultado.
En el siguiente ejemplo utilizaremos el mapa de profundidad de una imagen (la imagen original solo se usa para extraer el mapa de profundidad), junto con el prompt: dos astronautas bailando en un club ubicado en la luna.
Otros
ControlNet es una red neuronal que utiliza una serie de modelos para condicionar el resultado de maneras muy diversas. Es capaz de aplicar mapas de profundidad y poses entre otros muchos condicionamientos.
El siguiente ejemplo ilustra el condicionamiento de pose utilizando el prompt: astronauta en la luna.
No hace falta extraer la pose de una imagen mediante OpenPose, se puede crear desde cero. Por ejemplo, a través de este editor online.
El espacio latente
Falta una pieza del puzle que sin ella nada de esto podría realizarse. En todo este proceso se ha trabajado sobre los pixeles de la imagen, a esto se le conoce como espacio de imagen (image space). Con una imagen de 512x512 pixeles y colores RGB tendríamos un espacio de 786432 dimensiones (3x512x512), algo que no podrías ejecutar en el ordenador de tu casa ni en broma.
La solución a esto es utilizar el espacio latente, de ahí lo de modelo de difusión latente (LDM - latent diffusion model). En vez de trabajar sobre un espacio tan grande, primero se comprime a un espacio latente que es 64 veces más pequeño (3x64x64).
Te estarás preguntando como es posible comprimir la información de una imagen en algo 64 veces más pequeño sin perder detalle. Esto es posible porque no necesitamos toda la información. Es común que en la fotografía de una persona aparezca una cara, dos ojos, nariz, boca, etc. En vez de almacenar toda la información de una imagen se trabaja con características comunes y relevantes para representar la imagen de la manera más eficiente posible. Esta información se guarda en forma de números es un tensor de tamaño 3x64x64.
Un variational autoencoder (VAE) es un tipo de red neuronal que convierte una imagen en un tensor en el espacio latente (encoder) o un tensor del espacio latente en una imagen (decoder).
Todos los procesos que hemos visto en la técnica de difusión inversa tienen lugar en el espacio latente. Esto cambia el proceso en dos puntos:
- Al principio del proceso, en vez de generar una imagen llena de ruido se genera ruido latente y se guarda en un tensor. En el caso de inpainting, este tensor con ruido latente se genera pasando la imagen por el VAE encoder y después se aplica una transformación al tensor para añadir el ruido deseado.
- Al final del proceso, cuando Stable Diffusion tiene en el espacio latente el resultado final, el tensor pasa por el VAE decoder para convertirse en una imagen de 512x512 pixeles.
Stable Diffusion también te permite configurar que VAE utilizar. Dependiendo de la tarea unos ofrecen mejores resultados que otros, ya que en realidad sí se pierde algo de información en el espacio latente y el VAE es responsable de su recuperación.
Stable Diffusion y sus versiones
La versión 1.x de Stable Diffusion es la que hemos visto en este artículo. La diferencia del modelo 1.4 al modelo 1.5 ha sido el tiempo de entrenamiento, por lo que la versión 1.5 es capaz de generar imágenes de una mayor calidad.
En Stable Diffusion 2.x ha habido cambios más sustanciales. La versión 1.x utiliza el modelo CLIP ViT-L de Open AI mientras que Stable Diffusion 2.x utiliza el modelo OpenCLIP ViT-H de LAION. Este modelo es una versión de ViT-L/14 más grande y que ha sido entrenado con imágenes libres de derechos de autor. Se han utilizado imágenes de hasta 768x768 pixeles en contraposición a 1.x que utiliza imágenes de hasta 512x512, por lo que las imágenes generadas con la versión 2 tienen más calidad/resolución.
En cuanto a las imágenes libres de derechos de autor, es positivo para la comunidad pero también ha traído unas cuantas críticas. Los resultados en algunos casos son de peor calidad que en la versión anterior al generar imágenes utilizando estilos propietarios o caras de famosos. Tampoco se ha permitido la generacion de imágenes NSFW en esta versión.
En la versión 2.1 se ha relajado un poco el filtro NSFW durante el entrenamiento, por lo que ahora se pueden generar este tipo de imágenes.
La nueva versión SDXL 1.0 que saldrá a la luz en julio de 2023 será distribuida en 2 modelos: el modelo base y uno nuevo llamado refiner. En ambos se han utilizado imágenes con resolución de hasta 1024x1024 con diferentes relaciones de aspecto. El modelo base utiliza 2 modelos para el text encoding: OpenCLIP ViT-bigG de LAION y CLIP ViT-L de Open AI. En cuanto al modelo refiner solo utiliza el modelo OpenCLIP y además ha sido entrenado para afinar pequeños detalles, por lo que su uso está orientado a image-to-image. Ambos modelos ofrecen una calidad bastante más alta que sus versiones anteriores, además de mejor fidelidad en el prompt. A esto se le suma el hecho de que están trabajando con la comunidad para que puedan entrenar y afinar el modelo de manera óptima y así crear modelos derivados de muy alta calidad (como se conseguía con SD 1.5).
Conclusión
Para concluir vamos a recordar los pasos principales del proceso de inferencia en Stable Diffusion:
- Se genera un tensor con ruido latente en base a una semilla fija o aleatoria. En el caso de Inpainting se utiliza una imagen como base.
- Los condicionamientos (como texto, mapas de profundidad o etiquetas de clase) se convierten en vectores de embeddings que almacenan sus características en múltiples dimensiones.
- Estos vectores pasan por el transformer de CLIP que se encarga de calcular las relaciones entre los embeddings y sus características mediante la técnica de cross-attention.
- El predictor de ruido (U-Net) inicia el proceso de eliminación de ruido utilizando el resultado del transformer como guía hacia el resultado deseado. Este proceso utiliza la técnica de muestreo y se repite tantas veces como número de pasos se ha especificado.
- En este proceso, se genera ruido mediante el algoritmo de sampler seleccionado y se resta del tensor inicial. Este tensor más limpio de ruido sirve de base para el siguiente paso hasta completar todos los pasos. El noise scheduler controla la cantidad de progreso en cada paso para que no sea lineal.
- Cuando finaliza el proceso de eliminación de ruido, el tensor sale del espacio latente a través del decodificador del VAE, que se encarga de convertirlo en una imagen y así sorprendernos con el resultado.
Recuerda que puedes continuar en la segunda parte: Cómo implementar Stable Diffusion
Puedes apoyarme para que pueda dedicar aún más tiempo a escribir artículos y tener recursos para crear nuevos proyectos. ¡Gracias!