ControlNet

Os modelos ControlNet são adaptadores treinados sobre outro modelo pré-treinado. Isso permite um maior grau de controle sobre a geração de imagens ao condicionar o modelo com uma imagem de entrada adicional. A imagem de entrada pode ser uma borda canny, um mapa de profundidade, uma pose humana, entre muitas outras.

Se você estiver treinando em uma GPU com vRAM limitada, deve tentar habilitar os parâmetros gradient_checkpointing, gradient_accumulation_steps e mixed_precision no comando de treinamento. Você também pode reduzir o uso de memória utilizando atenção eficiente em memória com xFormers.

Este guia explorará o script de treinamento train_controlnet.py para ajudá-lo a se familiarizar com ele e como você pode adaptá-lo para o seu próprio caso de uso.

Antes de executar o script, certifique-se de instalar a biblioteca a partir do código-fonte:

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

Em seguida, navegue até a pasta de exemplo que contém o script de treinamento e instale as dependências necessárias para o script que você está usando:

cd examples/controlnet
pip install -r requirements.txt

🤗 Accelerate é uma biblioteca para ajudar você a treinar em múltiplos GPUs/TPUs ou com precisão mista. Ela configura automaticamente seu ambiente de treinamento com base no seu hardware e ambiente. Dê uma olhada no Quick tour do 🤗 Accelerate para saber mais.

Inicialize um ambiente do 🤗 Accelerate:

accelerate config

Para configurar um ambiente padrão do 🤗 Accelerate sem escolher nenhuma configuração:

accelerate config default

Ou, se o seu ambiente não suporta um shell interativo, como um notebook, você pode usar:

from accelerate.utils import write_basic_config

write_basic_config()

Por fim, se você deseja treinar um modelo em seu próprio conjunto de dados, consulte o guia Criar um conjunto de dados para treinamento para aprender a criar um conjunto de dados que funcione com o script de treinamento.

As seções a seguir destacam partes do script de treinamento que são importantes para entender como modificá-lo, mas não cobrem todos os aspectos do script em detalhes. Se você tiver interesse em aprender mais, sinta-se à vontade para ler o script e nos informar caso tenha alguma dúvida ou preocupação.

Parâmetros do script

O script de treinamento fornece muitos parâmetros para ajudar você a personalizar sua execução de treinamento. Todos os parâmetros e suas descrições estão encontrados na função parse_args(). Esta função fornece valores padrão para cada parâmetro, como o tamanho do lote de treinamento e a taxa de aprendizado, mas você também pode definir seus próprios valores no comando de treinamento, se desejar.

Por exemplo, para acelerar o treinamento com precisão mista usando o formato fp16, adicione o parâmetro --mixed_precision ao comando de treinamento:

accelerate launch train_controlnet.py \
  --mixed_precision="fp16"

Muitos dos parâmetros básicos e importantes são descritos no guia de treinamento Texto-para-imagem, portanto, este guia se concentra apenas nos parâmetros relevantes para o ControlNet:

--max_train_samples: o número de amostras de treinamento; esse valor pode ser reduzido para um treinamento mais rápido, mas, se você deseja transmitir conjuntos de dados realmente grandes, precisará incluir esse parâmetro e o parâmetro --streaming no seu comando de treinamento
--gradient_accumulation_steps: número de etapas de atualização a serem acumuladas antes da passagem reversa; isso permite treinar com um tamanho de lote maior do que a memória da sua GPU pode normalmente suportar

Ponderação Min-SNR

A estratégia de ponderação Min-SNR pode ajudar no treinamento ao reequilibrar a perda para alcançar uma convergência mais rápida. O script de treinamento suporta a previsão de epsilon (ruído) ou v_prediction, mas o Min-SNR é compatível com ambos os tipos de previsão. Essa estratégia de ponderação é suportada apenas pelo PyTorch.

Adicione o parâmetro --snr_gamma e defina-o com o valor recomendado de 5.0:

accelerate launch train_controlnet.py \
  --snr_gamma=5.0

Script de treinamento

Assim como com os parâmetros do script, um panorama geral do script de treinamento é fornecido no guia de treinamento Text-to-image. Em vez disso, este guia examina as partes relevantes do script do ControlNet.

O script de treinamento possui uma função make_train_dataset para pré-processar o conjunto de dados com transformações de imagem e tokenização das legendas. Você verá que, além da tokenização usual das legendas e das transformações de imagem, o script também inclui transformações para a imagem de condicionamento.

Se você estiver transmitindo um dataset em um TPU, o desempenho pode ser limitado pela biblioteca 🤗 Datasets, que não é otimizada para imagens. Para garantir a máxima taxa de transferência, recomendamos explorar outros formatos de dataset, como WebDataset, TorchData e TensorFlow Datasets.

conditioning_image_transforms = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution),
        transforms.ToTensor(),
    ]
)

Dentro da função main(), você encontrará o código para carregar o tokenizer, o text encoder, o scheduler e os modelos. É aqui que o modelo ControlNet é carregado, seja a partir de pesos existentes ou inicializado aleatoriamente a partir de um UNet:

if args.controlnet_model_name_or_path:
    logger.info("Loading existing controlnet weights")
    controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path)
else:
    logger.info("Initializing controlnet weights from unet")
    controlnet = ControlNetModel.from_unet(unet)

O optimizer é configurado para atualizar os parâmetros do ControlNet:

params_to_optimize = controlnet.parameters()
optimizer = optimizer_class(
    params_to_optimize,
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Por fim, no loop de treinamento, os embeddings do texto de condicionamento e a imagem são passados para os blocos down e mid do modelo ControlNet:

encoder_hidden_states = text_encoder(batch["input_ids"])[0]
controlnet_image = batch["conditioning_pixel_values"].to(dtype=weight_dtype)

down_block_res_samples, mid_block_res_sample = controlnet(
    noisy_latents,
    timesteps,
    encoder_hidden_states=encoder_hidden_states,
    controlnet_cond=controlnet_image,
    return_dict=False,
)

Se você quiser saber mais sobre como funciona o loop de treinamento, confira o tutorial Entendendo pipelines, modelos e agendadores, que detalha o padrão básico do processo de remoção de ruído.

Inicie o script

Agora você está pronto para iniciar o script de treinamento! 🚀

Este guia usa o conjunto de dados fusing/fill50k, mas lembre-se: você pode criar e usar seu próprio conjunto de dados, se desejar (veja o guia Criar um conjunto de dados para treinamento).

Defina a variável de ambiente MODEL_NAME para um id de modelo no Hub ou um caminho para um modelo local e OUTPUT_DIR para onde você deseja salvar o modelo.

Baixe as seguintes imagens para condicionar seu treinamento com:

wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_2.png

Mais uma coisa antes de você executar o script! Dependendo da GPU que você possui, pode ser necessário habilitar certas otimizações para treinar um ControlNet. A configuração padrão neste script requer ~38GB de vRAM. Se você estiver treinando em mais de uma GPU, adicione o parâmetro --multi_gpu ao comando accelerate launch.

16GB

12GB

8GB

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/save/model"

accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
 --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --push_to_hub

Quando o treinamento estiver concluído, você poderá usar o seu modelo recém-treinado para inferência!

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch

controlnet = ControlNetModel.from_pretrained("path/to/controlnet", dtype=torch.float16)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
    "path/to/base/model", controlnet=controlnet, dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

control_image = load_image("./conditioning_image_1.png")
prompt = "pale golden rod circle with old lace background"

generator = torch.manual_seed(0)
image = pipeline(prompt, num_inference_steps=20, generator=generator, image=control_image).images[0]
image.save("./output.png")

Stable Diffusion XL

Stable Diffusion XL (SDXL) é um modelo poderoso de texto para imagem que gera imagens em alta resolução e adiciona um segundo text-encoder à sua arquitetura. Use o script train_controlnet_sdxl.py para treinar um adaptador ControlNet para o modelo SDXL.

O script de treinamento do SDXL é discutido em mais detalhes no guia de treinamento do SDXL.

Próximas etapas

Parabéns por treinar seu próprio ControlNet! Para saber mais sobre como usar seu novo modelo, os seguintes guias podem ser úteis:

Aprenda como usar um ControlNet para inferência em uma variedade de tarefas.

Update on GitHub

Diffusers

ControlNet

Parâmetros do script

Ponderação Min-SNR

Script de treinamento

Inicie o script

Stable Diffusion XL

Próximas etapas