ControlNet

ControlNet es una arquitectura de red neuronal que añade control espacial a un modelo de difusión preentrenado de texto a imagen, como Stable Diffusion. Además del prompt de texto, recibe una imagen de condicionamiento, como un mapa de bordes, un mapa de profundidad, un esqueleto de pose humana o un mapa de segmentación; la imagen generada sigue esa estructura, mientras que el prompt determina el contenido y el estilo. Un prompt como «un caballero con armadura plateada», combinado con un esqueleto de pose, genera un caballero exactamente en esa pose. Lvmin Zhang, Anyi Rao y Maneesh Agrawala, de la Universidad de Stanford, lo presentaron en 2023 en Cómo añadir control condicional a los modelos de difusión de texto a imagen, artículo que compartió el premio Marr al mejor artículo en ICCV 2023. Los pesos del modelo base permanecen congelados, así que ControlNet añade control sin degradar lo que el modelo ya sabe.

Cómo funciona ControlNet#

Stable Diffusion genera una imagen mediante la eliminación repetida del ruido de un espacio latente con una U-Net. ControlNet bloquea esa U-Net y crea una copia entrenable de su codificador: los 12 bloques de codificación y el bloque central. Primero, una pequeña red convolucional convierte la imagen de condicionamiento de 512×512 en un mapa de características de 64×64 que coincide con el espacio latente. La copia entrenable procesa el espacio latente junto con estas características de condicionamiento, y sus salidas se añaden a las 12 conexiones de salto de la U-Net bloqueada y a su bloque central; así, el decodificador bloqueado recibe la señal de control:

La copia se conecta al modelo bloqueado únicamente mediante convoluciones cero: convoluciones de 1×1 cuyos pesos y sesgos empiezan en cero. En el primer paso de entrenamiento producen ceros, de modo que el modelo combinado se comporta exactamente igual que el original y ningún ruido aleatorio llega a las capas preentrenadas. La convolución cero de salida sigue recibiendo un gradiente de pesos distinto de cero, porque su entrada no es cero; por tanto, la primera actualización la aleja de cero y, a partir de ahí, aumenta la señal de control. Este código de PyTorch muestra un bloque:

import copy

import torch
from torch import nn

# A frozen block from the pretrained model and a trainable copy of it
torch.manual_seed(0)
locked = nn.Conv2d(64, 64, 3, padding=1).requires_grad_(False)
trainable = copy.deepcopy(locked).requires_grad_(True)

# Zero convolutions: 1x1 convolutions whose weights and biases start at zero
zero_in, zero_out = nn.Conv2d(64, 64, 1), nn.Conv2d(64, 64, 1)
for conv in (zero_in, zero_out):
    nn.init.zeros_(conv.weight)
    nn.init.zeros_(conv.bias)

x = torch.randn(1, 64, 32, 32)  # features entering the block
condition = torch.randn(1, 64, 32, 32)  # encoded edge, depth, or pose map
y = locked(x) + zero_out(trainable(x + zero_in(condition)))

# At step 0 the output equals the frozen block's, and only the outer zero conv gets a gradient
print(torch.equal(y, locked(x)))
y.sum().backward()
print(f"zero_in grad: {zero_in.weight.grad.norm():.1f}, zero_out grad: {zero_out.weight.grad.norm():.1f}")

Imprime True y zero_in grad: 0.0, zero_out grad: 2040.6. La salida del bloque es idéntica a la del modelo congelado, la convolución cero exterior recibe un gradiente porque su entrada no es cero y la interior empieza a aprender cuando la exterior deja de ser cero.

Tipos de condicionamiento#

El artículo original entrenó ControlNets independientes para bordes Canny, mapas de profundidad, mapas de normales de superficie, líneas rectas M-LSD, bordes suaves HED, mapas de segmentación de ADE20K, poses humanas de OpenPose y bocetos de usuarios. Cada imagen de condicionamiento se extrae de una foto de referencia con un modelo ya disponible, como el detector de bordes Canny de OpenCV o una red de estimación de profundidad, y se pueden aplicar varios ControlNet a la vez, cada uno con su propia intensidad. Los modelos de estimación de pose de Ultralytics YOLO26 extraen este tipo de condicionamiento de pose y predicen 17 puntos clave de COCO para cada persona, que se pueden dibujar como esqueletos sobre un lienzo en blanco:

import cv2
import numpy as np
from ultralytics import YOLO

# Carga un modelo de pose YOLO26 que prediga 17 puntos clave por persona
model = YOLO("yolo26n-pose.pt")
results = model("https://ultralytics.com/images/bus.jpg")
result = results[0]

# Dibuja solo los esqueletos sobre un lienzo negro del tamaño de la foto
skeletons = result.plot(img=np.zeros_like(result.orig_img), boxes=False)
cv2.imwrite("pose_condition.png", skeletons)

La imagen guardada muestra las extremidades de cada persona como líneas de colores sobre un fondo negro, sin los colores, las texturas ni el fondo de la foto.

ControlNet frente a métodos relacionados#

• LoRA: Ajusta un pequeño conjunto de pesos añadidos para enseñar al modelo un estilo, personaje u objeto nuevo. LoRA cambia lo que dibuja el modelo, mientras que ControlNet fija dónde aparecen los elementos; a menudo se combinan ambos.
• T2I-Adapter: Alternativa más ligera, con unos 77 millones de parámetros, que añade características de condicionamiento a las características del codificador de la U-Net congelada en lugar de ejecutar una copia completa del codificador. Los adaptadores para distintas condiciones se pueden combinar durante la generación.
• Generación de imagen a imagen: Empieza a eliminar el ruido de una copia de la foto de entrada a la que se ha añadido ruido, por lo que la salida hereda sus colores y su composición aproximada. ControlNet limita únicamente la estructura extraída, como los bordes, y deja los colores y las texturas en manos del prompt.
• Relleno de imágenes: Regenera una parte enmascarada de una imagen para que encaje con el resto. Controla qué región cambia, mientras que ControlNet controla la estructura de lo que se genera.

Aplicaciones en el mundo real#

• Arquitectura y diseño de interiores: Los bocetos, los dibujos lineales CAD o los mapas de líneas M-LSD de una habitación se convierten en representaciones fotorrealistas con distintos materiales y estilos, manteniendo la misma geometría.
• Arte de personajes y animación: Los esqueletos de pose mantienen fija la postura de un personaje en distintas variaciones de estilo, y las poses de cada fotograma extraídas de un vídeo de referencia permiten crear secuencias animadas coherentes.
• Imágenes de productos y marketing: Los mapas de profundidad o de bordes de una foto de producto mantienen su forma y ubicación, mientras el prompt cambia el entorno, la iluminación y el fondo.
• Datos sintéticos: Generar imágenes a partir de mapas de segmentación o poses produce nuevas imágenes de entrenamiento cuyas etiquetas se conocen de antemano, porque la composición se especificó antes de la generación.

Entrenamiento de ControlNet#

El entrenamiento requiere pares de imágenes y sus mapas de condicionamiento, que se generan automáticamente ejecutando el extractor de condiciones sobre un conjunto de datos de imágenes con descripciones. Los pesos bloqueados no necesitan actualizaciones de gradiente, por lo que cada iteración de entrenamiento requiere solo alrededor de un 23 % más de memoria de GPU y un 34 % más de tiempo que optimizar Stable Diffusion sin ControlNet; el artículo informa de un entrenamiento sólido con conjuntos de datos de menos de 50 000 a más de un millón de imágenes. El entrenamiento también muestra una «convergencia repentina»: el modelo ignora el condicionamiento durante miles de pasos y luego empieza a seguirlo bruscamente, normalmente antes de completar 10 000 pasos de optimización; en el ejemplo del artículo, ocurre en el paso 6133.