ControlNet

Модели ControlNet представляют собой адаптеры, обученные поверх другой предобученной модели. Это позволяет более точно управлять генерацией изображений, используя дополнительное входное изображение в качестве условия. В качестве входного изображения могут выступать канни-контур, карта глубины, поза человека и многое другое.

Если вы обучаете модель на GPU с ограниченным объёмом vRAM, рекомендуется включить параметры gradient_checkpointing, gradient_accumulation_steps и mixed_precision в команде обучения. Кроме того, вы можете снизить потребление памяти, используя энергоэффективное внимание с xFormers.

В этом руководстве мы рассмотрим скрипт обучения train_controlnet.py, чтобы вы могли ознакомиться с ним и понять, как адаптировать его под свои задачи.

Перед запуском скрипта убедитесь, что вы установили библиотеку из исходного кода:

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

Затем перейдите в папку с примером, содержащую скрипт обучения, и установите необходимые зависимости для используемого вами скрипта:

cd examples/controlnet
pip install -r requirements.txt

🤗 Accelerate — это библиотека, которая помогает обучать модели на нескольких GPU/TPU или с использованием смешанной точности. Она автоматически настраивает среду обучения в зависимости от вашего оборудования и окружения. Подробнее можно узнать в Кратком обзоре 🤗 Accelerate.

Инициализируйте среду 🤗 Accelerate:

accelerate config

Чтобы настроить среду по умолчанию 🤗 Accelerate без выбора каких-либо конфигураций:

accelerate config default

Или, если ваша среда не поддерживает интерактивную оболочку, например, ноутбук, вы можете использовать:

from accelerate.utils import write_basic_config

write_basic_config()

Наконец, если вы хотите обучить модель на собственном наборе данных, ознакомьтесь с руководством Create a dataset for training, чтобы узнать, как создать набор данных, совместимый со скриптом обучения.

В следующих разделах выделены части скрипта обучения, важные для понимания того, как его модифицировать, но не все аспекты скрипта описаны подробно. Если вы хотите узнать больше, не стесняйтесь прочитать скрипт и сообщить нам, если у вас возникнут вопросы или замечания.

Параметры скрипта

Скрипт обучения предоставляет множество параметров, которые помогают настроить процесс обучения. Все параметры и их описания находятся в функции parse_args(). Эта функция задаёт значения по умолчанию для каждого параметра, например, размер обучающей выборки и скорость обучения, но при необходимости вы можете указать свои значения в команде запуска обучения.

Например, чтобы ускорить обучение с использованием смешанной точности в формате fp16, добавьте параметр --mixed_precision к команде обучения:

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

Многие базовые и важные параметры описаны в руководстве по обучению Text-to-image, поэтому в данном руководстве рассматриваются только параметры, относящиеся к ControlNet:

• --max_train_samples: количество обучающих примеров; это значение можно уменьшить для ускорения обучения, но если вы хотите стримить действительно большие наборы данных, вам нужно включить в команду обучения этот параметр и параметр --streaming
• --gradient_accumulation_steps: количество шагов обновления, которые накапливаются перед обратным распространением; это позволяет обучать модель с размером батча, превышающим то, что обычно может обработать память вашего GPU

Взвешивание Min-SNR

Стратегия взвешивания Min-SNR может помочь в обучении, перераспределяя функцию потерь для достижения более быстрой сходимости. Скрипт обучения поддерживает предсказание epsilon (шума) или v_prediction, при этом Min-SNR совместима с обоими типами предсказаний. Данная стратегия взвешивания поддерживается только в PyTorch.

Добавьте параметр --snr_gamma и установите для него рекомендуемое значение 5.0:

accelerate launch train_controlnet.py \
  --snr_gamma=5.0

Скрипт обучения

Как и в случае с параметрами скрипта, общий обзор скрипта обучения приведён в руководстве по обучению Text-to-image. В данном руководстве рассматриваются только соответствующие части скрипта ControlNet.

Скрипт обучения содержит функцию make_train_dataset для предобработки набора данных с применением трансформаций изображений и токенизации подписей. Вы увидите, что помимо стандартной токенизации подписей и трансформаций изображений, скрипт также включает трансформации для условного изображения.

Если вы используете стриминг набора данных на TPU, производительность может ограничиваться библиотекой 🤗 Datasets, которая не оптимизирована для работы с изображениями. Для обеспечения максимальной пропускной способности рекомендуется рассмотреть другие форматы наборов данных, такие как WebDataset, TorchData и TensorFlow Datasets.

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

В функции main() находится код для загрузки токенизатора, текстового энкодера, планировщика и моделей. Здесь же загружается модель ControlNet — либо из существующих весов, либо с случайной инициализацией из 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)

Оптимизатор настроен для обновления параметров 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,
)

Наконец, в цикле обучения эмбеддинги текста-условия и изображение передаются в нижние и средние блоки модели 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,
)

Если вы хотите узнать больше о том, как работает цикл обучения, обратитесь к учебнику Understanding pipelines, models and schedulers, в котором подробно разбирается базовый паттерн процесса денойзинга.

Запустите скрипт

Теперь вы готовы запустить скрипт обучения! 🚀

В этом руководстве используется набор данных fusing/fill50k, но помните, что при желании вы можете создать и использовать собственный набор данных (см. руководство Создание набора данных для обучения).

Установите переменную окружения MODEL_NAME в идентификатор модели на Hub или путь к локальной модели, а OUTPUT_DIR — в каталог, куда вы хотите сохранить модель.

Скачайте следующие изображения, чтобы использовать их для настройки обучения:

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

Ещё одно замечание перед запуском скрипта! В зависимости от вашей GPU может потребоваться включить определённые оптимизации для обучения ControlNet. Конфигурация по умолчанию в этом скрипте требует ~38 ГБ vRAM. Если вы обучаете на нескольких GPU, добавьте параметр --multi_gpu к команде accelerate launch.

pip install bitsandbytes

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

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

accelerate config

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

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

После завершения обучения вы можете использовать вашу новую модель для инференса!

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) — это мощная модель text-to-image, генерирующая изображения высокого разрешения, которая добавляет второй text-encoder в свою архитектуру. Используйте скрипт train_controlnet_sdxl.py для обучения адаптера ControlNet для модели SDXL.

Скрипт обучения SDXL подробно описан в руководстве обучение SDXL.

Дальнейшие шаги

Поздравляем с обучением собственной модели ControlNet! Чтобы узнать, как использовать новую модель, вам могут помочь следующие руководства:

• Узнайте, как использовать ControlNet для инференса в различных задачах.