Custom Diffusion

Custom Diffusion — это метод обучения для персонализации моделей генерации изображений. Как и Textual Inversion, DreamBooth и LoRA, Custom Diffusion требует лишь нескольких (~4–5) примеров изображений. Этот метод работает за счёт обучения только весов в слоях cross-attention и использует специальное слово для представления нового концепта. Особенность Custom Diffusion в том, что он способен обучаться нескольким концептам одновременно.

Если вы обучаете модель на GPU с ограниченным объёмом vRAM, рекомендуется включить xFormers с помощью --enable_xformers_memory_efficient_attention для ускорения обучения и снижения требований к vRAM (16 ГБ). Чтобы сэкономить ещё больше памяти, добавьте --set_grads_to_none в аргументы обучения, чтобы устанавливать градиенты в None вместо нуля (этот параметр может вызывать некоторые проблемы, поэтому, если вы с ними столкнётесь, попробуйте его удалить).

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

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

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

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

cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval

🤗 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(). Функция имеет значения по умолчанию, но при необходимости вы можете задать свои значения в команде запуска обучения.

Например, чтобы изменить разрешение входного изображения:

accelerate launch train_custom_diffusion.py \
  --resolution=256

Многие базовые параметры описаны в руководстве по обучению DreamBooth, поэтому в этом руководстве мы сосредоточимся на параметрах, уникальных для Custom Diffusion:

--freeze_model: замораживает параметры ключей и значений в слое кросс-внимания; по умолчанию установлено значение crossattn_kv, но вы можете задать crossattn, чтобы обучать все параметры в слое кросс-внимания
--concepts_list: чтобы обучить несколько концепций, укажите путь к JSON-файлу, содержащему эти концепции
--modifier_token: специальное слово, используемое для представления изученного концепта
--initializer_token: специальное слово, используемое для инициализации эмбеддингов modifier_token

Потеря сохранения априорных знаний

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

Многие параметры для потерь сохранения предыдущего состояния описаны в руководстве по обучению DreamBooth.

Регуляризация

Custom Diffusion включает обучение целевых изображений с использованием небольшого набора реальных изображений для предотвращения переобучения. Как вы можете представить, это легко сделать, когда обучение выполняется всего на нескольких изображениях! Скачайте 200 реальных изображений с помощью clip_retrieval. Значение class_prompt должно соответствовать той же категории, что и целевые изображения. Эти изображения сохраняются в class_data_dir.

python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200

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

--with_prior_preservation: использовать ли потерю сохранения априорного распределения
--prior_loss_weight: управляет влиянием функции потерь сохранения априорного распределения на модель
--real_prior: использовать ли небольшой набор реальных изображений для предотвращения переобучения

accelerate launch train_custom_diffusion.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="./real_reg/samples_cat" \
  --class_prompt="cat" \
  --real_prior=True \

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

Многие фрагменты кода в скрипте обучения Custom Diffusion похожи на код скрипта DreamBooth. В этом руководстве мы сосредоточимся на коде, относящемся именно к Custom Diffusion.

Скрипт обучения Custom Diffusion содержит два класса датасетов:

CustomDiffusionDataset: предобработка изображений, класс-изображений и промптов для обучения
PromptDataset: готовит промпты для генерации изображений классов

Далее modifier_token добавляется в токенизатор, преобразуется в идентификаторы токенов, а размеры эмбеддингов токенов изменяются с учётом нового modifier_token. Затем эмбеддинги modifier_token инициализируются эмбеддингами initializer_token. Все параметры текстового энкодера замораживаются, за исключением эмбеддингов токенов, поскольку именно с ними модель пытается ассоциировать концепты.

params_to_freeze = itertools.chain(
    text_encoder.text_model.encoder.parameters(),
    text_encoder.text_model.final_layer_norm.parameters(),
    text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)

Теперь вам нужно добавить Custom Diffusion weights в слои внимания. Это очень важный шаг для корректного определения формы и размера весов внимания, а также для установки нужного количества процессоров внимания в каждом блоке UNet.

st = unet.state_dict()
for name, _ in unet.attn_processors.items():
    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
    if name.startswith("mid_block"):
        hidden_size = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"):
        block_id = int(name[len("up_blocks.")])
        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
    elif name.startswith("down_blocks"):
        block_id = int(name[len("down_blocks.")])
        hidden_size = unet.config.block_out_channels[block_id]
    layer_name = name.split(".processor")[0]
    weights = {
        "to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
        "to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
    }
    if train_q_out:
        weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
        weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
        weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
    if cross_attention_dim is not None:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=train_kv,
            train_q_out=train_q_out,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        ).to(unet.device)
        custom_diffusion_attn_procs[name].load_state_dict(weights)
    else:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=False,
            train_q_out=False,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        )
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)

optimizer инициализируется для обновления параметров слоя кросс-внимания:

optimizer = optimizer_class(
    itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
    if args.modifier_token is not None
    else custom_diffusion_layers.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

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

if args.modifier_token is not None:
    if accelerator.num_processes > 1:
        grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
    else:
        grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
    index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
    for i in range(len(modifier_token_id[1:])):
        index_grads_to_zero = index_grads_to_zero & (
            torch.arange(len(tokenizer)) != modifier_token_id[i]
        )
    grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
        index_grads_to_zero, :
    ].fill_(0)

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

Когда вы внесёте все изменения или решите оставить конфигурацию по умолчанию, можно запускать скрипт обучения! 🚀

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

Установите переменную окружения MODEL_NAME в идентификатор модели на Hub или путь к локальной модели, INSTANCE_DIR — в путь, куда вы только что скачали изображения с котами, и OUTPUT_DIR — в место, где вы хотите сохранить модель. В качестве специального слова, привязывающего вновь обученные эмбеддинги, используйте <new1>. Скрипт создаёт и сохраняет контрольные точки модели и файл pytorch_custom_diffusion_weights.bin в вашем репозитории.

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

Если вы обучаете модель на изображениях человеческих лиц, команда Custom Diffusion обнаружила, что следующие параметры работают хорошо:

--learning_rate=5e-6

--max_train_steps может принимать любое значение от 1000 до 2000

--freeze_model=crossattn

используйте не менее 15–20 изображений для обучения

одно понятие

несколько концепций

Когда обучение завершено, вы можете использовать вашу новую модель Custom Diffusion для инференса.

одно понятие

несколько концепций

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

Поздравляем с обучением модели с помощью Custom Diffusion! 🎉 Чтобы узнать больше:

Ознакомьтесь с статьей в блоге о многоконцептуальной кастомизации текстово-изображенийых диффузионных моделей, чтобы узнать больше о результатах экспериментов команды Custom Diffusion.

Update on GitHub

Diffusers

Custom Diffusion

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

Потеря сохранения априорных знаний

Регуляризация

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

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

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