# ControlNet

Os modelos [ControlNet](https://hf.co/papers/2302.05543) 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](../optimization/attention_backends).

Este guia explorará o script de treinamento [train_controlnet.py](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/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:

```bash
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:

```bash
cd examples/controlnet
pip install -r requirements.txt
```

> [!TIP]
> 🤗 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](https://huggingface.co/docs/accelerate/quicktour) do 🤗 Accelerate para saber mais.

Inicialize um ambiente do 🤗 Accelerate:

```bash
accelerate config
```

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

```bash
accelerate config default
```

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

```py
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](create_dataset) para aprender a criar um conjunto de dados que funcione com o script de treinamento.

> [!TIP]
> 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](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet.py) 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()`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L231). 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:

```bash
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](text2image#script-parameters), 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](https://huggingface.co/papers/2303.09556) 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:

```bash
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](text2image#training-script). Em vez disso, este guia examina as partes relevantes do script do ControlNet.

O script de treinamento possui uma função [`make_train_dataset`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L582) 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.

> [!TIP]
> 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](https://github.com/webdataset/webdataset), [TorchData](https://github.com/pytorch/data) e [TensorFlow Datasets](https://www.tensorflow.org/datasets/tfless_tfds).

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

Dentro da função [`main()`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L713), 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:

```py
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](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L871) é configurado para atualizar os parâmetros do ControlNet:

```py
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](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/controlnet/train_controlnet.py#L943), os embeddings do texto de condicionamento e a imagem são passados para os blocos down e mid do modelo ControlNet:

```py
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](../using-diffusers/write_own_pipeline), 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](https://huggingface.co/datasets/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](create_dataset)).

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:

```bash
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`.

Em uma GPU de 16GB, você pode usar o otimizador de 8 bits do bitsandbytes e o gradient checkpointing para otimizar sua execução de treinamento. Instale o bitsandbytes:

```py
pip install bitsandbytes
```

Em seguida, adicione o seguinte parâmetro ao seu comando de treinamento:

```bash
accelerate launch train_controlnet.py \
  --gradient_checkpointing \
  --use_8bit_adam \
```

Em uma GPU de 12GB, você precisará do otimizador de 8 bits do bitsandbytes, do gradient checkpointing, do xFormers e definir os gradientes como `None` em vez de zero para reduzir o uso de memória.

```bash
accelerate launch train_controlnet.py \
  --use_8bit_adam \
  --gradient_checkpointing \
  --enable_xformers_memory_efficient_attention \
  --set_grads_to_none \
```

Em uma GPU de 8 GB, você precisará usar [DeepSpeed](https://www.deepspeed.ai/) para transferir alguns dos tensores da vRAM para a CPU ou o NVME, permitindo o treinamento com menos memória de GPU.

Execute o seguinte comando para configurar seu ambiente 🤗 Accelerate:

```bash
accelerate config
```

Durante a configuração, confirme que deseja usar o DeepSpeed stage 2. Agora, deve ser possível treinar com menos de 8GB de vRAM combinando o DeepSpeed stage 2, a precisão mista fp16 e o offloading dos parâmetros do modelo e do estado do otimizador para a CPU. A desvantagem é que isso requer mais RAM de sistema (~25 GB). Consulte a [documentação do DeepSpeed](https://huggingface.co/docs/accelerate/usage_guides/deepspeed) para mais opções de configuração. Seu arquivo de configuração deve se parecer com:

```bash
compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 4
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: false
  zero_stage: 2
distributed_type: DEEPSPEED
```

Você também deve alterar o otimizador Adam padrão para a versão otimizada do Adam do DeepSpeed [`deepspeed.ops.adam.DeepSpeedCPUAdam`](https://deepspeed.readthedocs.io/en/latest/optimizers.html#adam-cpu) para uma aceleração substancial. A ativação de `DeepSpeedCPUAdam` exige que a versão da toolchain CUDA do seu sistema seja a mesma da instalada com o PyTorch.

Os otimizadores de 8 bits do bitsandbytes não parecem ser compatíveis com o DeepSpeed no momento.

É isso! Você não precisa adicionar nenhum parâmetro adicional ao seu comando de treinamento.

```bash
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!

```py
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`](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/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](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](../using-diffusers/controlnet) para inferência em uma variedade de tarefas.