T2I-Adapter

T2I-Adapter — это лёгкая адаптерная модель, которая обеспечивает дополнительный входной образ для условного управления (line art, canny, sketch, depth, pose) для более точного контроля генерации изображений. Она похожа на ControlNet, но значительно меньше (~77M параметров и ~300MB размер файла), поскольку вставляет веса непосредственно в UNet, а не копирует и обучает его.

T2I-Adapter доступен только для обучения с моделью Stable Diffusion XL (SDXL).

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

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

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

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

cd examples/t2i_adapter
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(). Для каждого параметра заданы значения по умолчанию, например, размер обучающей выборки и скорость обучения, но при необходимости вы можете указать свои значения в команде запуска обучения.

Например, чтобы включить накопление градиентов, добавьте параметр --gradient_accumulation_steps к команде обучения:

accelerate launch train_t2i_adapter_sdxl.py \
  ----gradient_accumulation_steps=4

Many of the basic and important parameters are described in the Text-to-image training guide, so this guide just focuses on the relevant T2I-Adapter parameters:

• --pretrained_vae_model_name_or_path: путь к предобученной VAE; известно, что VAE в SDXL страдает от численной нестабильности, поэтому этот параметр позволяет указать более подходящую VAE
• --crops_coords_top_left_h и --crops_coords_top_left_w: координаты высоты и ширины, включаемые в эмбеддинги координат обрезки SDXL
• --conditioning_image_column: столбец с изображениями для условного формирования в наборе данных
• --proportion_empty_prompts: доля текстовых подсказок для изображений, которые заменяются пустыми строками

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

As with the script parameters, a walkthrough of the training script is provided in the Text-to-image training guide. Instead, this guide takes a look at the T2I-Adapter relevant parts of the script.

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

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

В функции main() T2I-Adapter загружается из предобученного адаптера или инициализируется случайным образом:

if args.adapter_model_name_or_path:
    logger.info("Loading existing adapter weights.")
    t2iadapter = T2IAdapter.from_pretrained(args.adapter_model_name_or_path)
else:
    logger.info("Initializing t2iadapter weights.")
    t2iadapter = T2IAdapter(
        in_channels=3,
        channels=(320, 640, 1280, 1280),
        num_res_blocks=2,
        downscale_factor=16,
        adapter_type="full_adapter_xl",
    )

Оптимизатор инициализируется для параметров T2I-Adapter:

params_to_optimize = t2iadapter.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,
)

Наконец, в цикле обучения изображение-условие адаптера и текстовые эмбеддинги передаются в UNet для предсказания шумового остатка:

t2iadapter_image = batch["conditioning_pixel_values"].to(dtype=weight_dtype)
down_block_additional_residuals = t2iadapter(t2iadapter_image)
down_block_additional_residuals = [
    sample.to(dtype=weight_dtype) for sample in down_block_additional_residuals
]

model_pred = unet(
    inp_noisy_latents,
    timesteps,
    encoder_hidden_states=batch["prompt_ids"],
    added_cond_kwargs=batch["unet_added_conditions"],
    down_block_additional_residuals=down_block_additional_residuals,
).sample

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

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

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

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

Установите переменную окружения MODEL_DIR в идентификатор модели на 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

Для отслеживания хода обучения с помощью Weights & Biases добавьте параметр --report_to=wandb к команде обучения. Также необходимо добавить в команду обучения параметры --validation_image, --validation_prompt и --validation_steps, чтобы фиксировать результаты. Это очень полезно для отладки модели и просмотра промежуточных результатов.

export MODEL_DIR="stabilityai/stable-diffusion-xl-base-1.0"
export OUTPUT_DIR="path to save model"

accelerate launch train_t2i_adapter_sdxl.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --mixed_precision="fp16" \
 --resolution=1024 \
 --learning_rate=1e-5 \
 --max_train_steps=15000 \
 --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
 --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
 --validation_steps=100 \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --report_to="wandb" \
 --seed=42 \
 --push_to_hub

После завершения обучения вы можете использовать свой T2I-Adapter для инференса:

from diffusers import StableDiffusionXLAdapterPipeline, T2IAdapter, EulerAncestralDiscreteSchedulerTest
from diffusers.utils import load_image
import torch

adapter = T2IAdapter.from_pretrained("path/to/adapter", dtype=torch.float16)
pipeline = StableDiffusionXLAdapterPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", adapter=adapter, dtype=torch.float16
)

pipeline.scheduler = EulerAncestralDiscreteSchedulerTest.from_config(pipe.scheduler.config)
pipeline.enable_xformers_memory_efficient_attention()
pipeline.enable_model_cpu_offload()

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, image=control_image, generator=generator
).images[0]
image.save("./output.png")

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

Поздравляем с обучением модели T2I-Adapter! 🎉 Чтобы узнать больше:

• Ознакомьтесь с статьей в блоге «Эффективная управляемая генерация для SDXL с помощью T2I-Adapters», чтобы узнать больше о результатах экспериментов команды T2I-Adapter.