ControlNet

ControlNet é uma arquitetura de rede neuronal que acrescenta controlo espacial a um modelo de difusão de texto para imagem pré-treinado, como o Stable Diffusion. Para além do prompt de texto, recebe uma imagem de condicionamento, como um mapa de contornos, um mapa de profundidade, um esqueleto de pose humana ou um mapa de segmentação; a imagem gerada segue essa estrutura, enquanto o prompt determina o conteúdo e o estilo. Um prompt como «um cavaleiro com armadura prateada», combinado com um esqueleto de pose, produz um cavaleiro exatamente nessa pose. Lvmin Zhang, Anyi Rao e Maneesh Agrawala, da Universidade Stanford, apresentaram-no em 2023 no artigo Adding Conditional Control to Text-to-Image Diffusion Models, que partilhou o prémio Marr para melhor artigo na ICCV 2023. Os pesos do modelo base permanecem congelados, pelo que o ControlNet acrescenta controlo sem degradar o que o modelo já sabe.

Como funciona o ControlNet#

O Stable Diffusion gera uma imagem ao remover repetidamente o ruído de uma representação latente com uma U-Net. O ControlNet bloqueia essa U-Net e cria uma cópia treinável do respetivo codificador: os 12 blocos de codificação e o bloco intermédio. Uma pequena rede convolucional converte primeiro a imagem de condicionamento de 512×512 num mapa de características de 64×64 que corresponde à representação latente. A cópia treinável processa a representação latente e estas características de condicionamento, e as respetivas saídas são adicionadas às 12 ligações de salto da U-Net bloqueada e ao seu bloco intermédio, para que o descodificador bloqueado receba o sinal de controlo:

A cópia liga-se ao modelo bloqueado apenas através de convoluções zero: convoluções 1×1 cujos pesos e desvios começam a zero. No primeiro passo de treino, produzem zeros, pelo que o modelo combinado se comporta exatamente como o original e nenhum ruído aleatório chega às camadas pré-treinadas. A convolução zero de saída continua a receber um gradiente de peso diferente de zero, porque a sua entrada não é zero; assim, a primeira atualização afasta-a de zero e o sinal de controlo começa a crescer. Este código PyTorch mostra um bloco:

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 e zero_in grad: 0.0, zero_out grad: 2040.6. A saída do bloco é idêntica à do modelo congelado; a convolução zero exterior recebe um gradiente porque a sua entrada é diferente de zero, e a convolução interior começa a aprender quando a exterior deixa de ser zero.

Tipos de condicionamento#

O artigo original treinou ControlNets separados para contornos Canny, mapas de profundidade, mapas de normais de superfície, linhas retas M-LSD, contornos suaves HED, mapas de segmentação ADE20K, poses humanas OpenPose e esboços de utilizadores. Cada imagem de condicionamento é extraída de uma fotografia de referência com um modelo pronto a utilizar, como o detetor de contornos Canny do OpenCV ou uma rede de estimação de profundidade; é possível aplicar vários ControlNets em simultâneo, cada um com a sua própria intensidade. Os modelos de estimação de pose Ultralytics YOLO26 extraem este tipo de condicionamento de pose, prevendo 17 pontos-chave COCO para cada pessoa, que podem ser desenhados como esqueletos num fundo em branco:

import cv2
import numpy as np
from ultralytics import YOLO

# Carrega um modelo de pose YOLO26 que prevê 17 pontos-chave por pessoa
model = YOLO("yolo26n-pose.pt")
results = model("https://ultralytics.com/images/bus.jpg")
result = results[0]

# Desenha apenas os esqueletos, num fundo preto do tamanho da fotografia
skeletons = result.plot(img=np.zeros_like(result.orig_img), boxes=False)
cv2.imwrite("pose_condition.png", skeletons)

A imagem guardada mostra os membros de cada pessoa como linhas coloridas sobre um fundo preto, sem as cores, texturas ou o fundo da fotografia.

ControlNet e métodos relacionados#

• LoRA: Ajusta um pequeno conjunto de pesos acrescentados para ensinar ao modelo um novo estilo, personagem ou objeto. A LoRA altera aquilo que o modelo desenha, enquanto o ControlNet fixa a posição dos elementos; as duas técnicas são frequentemente combinadas.
• T2I-Adapter: Alternativa mais leve, com cerca de 77 milhões de parâmetros, que acrescenta características de condicionamento às características do codificador da U-Net congelada, em vez de executar uma cópia completa do codificador. É possível combinar adaptadores para diferentes tipos de condicionamento durante a geração.
• Geração de imagem para imagem: Começa por remover o ruído de uma cópia com ruído da fotografia de entrada, pelo que a saída herda as cores e a disposição geral dessa fotografia. O ControlNet restringe apenas a estrutura extraída, como os contornos, e deixa as cores e as texturas a cargo do prompt.
• Preenchimento de imagens: Regenera uma parte mascarada de uma imagem para a integrar no restante conteúdo. Controla a região que muda, enquanto o ControlNet controla a estrutura do conteúdo gerado.

Aplicações no mundo real#

• Arquitetura e design de interiores: Esboços, desenhos de linhas CAD ou mapas de linhas M-LSD de uma divisão transformam-se em imagens fotorrealistas com materiais e estilos diferentes, mantendo a mesma geometria.
• Arte e animação de personagens: Os esqueletos de pose mantêm fixa a postura de uma personagem entre variações de estilo, e as poses por fotograma extraídas de um vídeo de referência orientam sequências animadas consistentes.
• Imagens de produtos e marketing: Os mapas de profundidade ou de contornos de uma fotografia de produto preservam a forma e a posição do produto, enquanto o prompt altera o cenário, a iluminação e o fundo.
• Dados sintéticos: A geração de imagens a partir de mapas de segmentação ou poses produz novas imagens de treino com etiquetas conhecidas antecipadamente, porque a disposição foi especificada antes da geração.

Treinar um ControlNet#

O treino requer pares de imagens e mapas de condicionamento, produzidos automaticamente ao executar o extrator de condicionamento num conjunto de dados de imagens legendadas. Os pesos bloqueados não precisam de atualizações de gradiente, pelo que cada iteração de treino requer apenas cerca de 23% mais memória da GPU e 34% mais tempo do que otimizar o Stable Diffusion sem ControlNet; o artigo relata um treino robusto em conjuntos de dados com menos de 50 000 e mais de um milhão de imagens. O treino também revela uma «convergência súbita»: o modelo ignora o condicionamento durante milhares de passos e, depois, começa a segui-lo abruptamente, geralmente em menos de 10 000 passos de otimização; no exemplo do artigo, isso acontece no passo 6133.