Buckets:

hf-doc-build
/

doc

Files
xet
hf-doc-build/doc / diffusers /main /ru /training /custom_diffusion.md
mishig
8 days ago
preview code
download
raw
22.1 kB
# Custom Diffusion

[Custom Diffusion](https://huggingface.co/papers/2212.04488) — это метод обучения для персонализации моделей генерации изображений. Как и 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](https://github.com/huggingface/diffusers/blob/main/examples/custom_diffusion/train_custom_diffusion.py), чтобы вы лучше с ним познакомились и могли адаптировать его под свои задачи.

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

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

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

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

> [!TIP]
> 🤗 Accelerate — это библиотека, которая помогает обучать модели на нескольких GPU/TPU или с использованием смешанной точности. Она автоматически настраивает среду обучения в зависимости от вашего оборудования и окружения. Подробнее можно узнать в [Кратком обзоре](https://huggingface.co/docs/accelerate/quicktour) 🤗 Accelerate.

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

```bash
accelerate config
```

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

```bash
accelerate config default
```

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

```py
from accelerate.utils import write_basic_config
write_basic_config()
```

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

> [!TIP]
> В следующих разделах выделены части скрипта обучения, важные для понимания того, как его модифицировать, но не все аспекты скрипта описаны подробно. Если вы хотите узнать больше, не стесняйтесь прочитать [скрипт](https://github.com/huggingface/diffusers/blob/main/examples/custom_diffusion/train_custom_diffusion.py) и сообщить нам, если у вас возникнут вопросы или замечания.

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

Скрипт обучения содержит все параметры, которые помогают настроить процесс обучения. Они находятся в функции [`parse_args()`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L319). Функция имеет значения по умолчанию, но при необходимости вы можете задать свои значения в команде запуска обучения.

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

```bash
accelerate launch train_custom_diffusion.py \
  --resolution=256
```

Многие базовые параметры описаны в руководстве по обучению [DreamBooth](dreambooth#script-parameters), поэтому в этом руководстве мы сосредоточимся на параметрах, уникальных для Custom Diffusion:

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

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

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

Многие параметры для потерь сохранения предыдущего состояния описаны в руководстве по обучению [DreamBooth](dreambooth#prior-preservation-loss).

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

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

```bash
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`: использовать ли небольшой набор реальных изображений для предотвращения переобучения

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

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

> [!TIP]
> Многие фрагменты кода в скрипте обучения Custom Diffusion похожи на код скрипта [DreamBooth](dreambooth#training-script). В этом руководстве мы сосредоточимся на коде, относящемся именно к Custom Diffusion.

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

- [`CustomDiffusionDataset`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L165): предобработка изображений, класс-изображений и промптов для обучения
- [`PromptDataset`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L148): готовит промпты для генерации изображений классов

Далее `modifier_token` [добавляется в токенизатор](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L811), преобразуется в идентификаторы токенов, а размеры эмбеддингов токенов изменяются с учётом нового `modifier_token`. Затем эмбеддинги `modifier_token` инициализируются эмбеддингами `initializer_token`. Все параметры текстового энкодера замораживаются, за исключением эмбеддингов токенов, поскольку именно с ними модель пытается ассоциировать концепты.

```py
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](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L911C3-L911C3) в слои внимания. Это очень важный шаг для корректного определения формы и размера весов внимания, а также для установки нужного количества процессоров внимания в каждом блоке UNet.

```py
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](https://github.com/huggingface/diffusers/blob/84cd9e8d01adb47f046b1ee449fc76a0c32dc4e2/examples/custom_diffusion/train_custom_diffusion.py#L982) инициализируется для обновления параметров слоя кросс-внимания:

```py
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,
)
```

В [цикле обучения](https://github.com/huggingface/diffusers/blob/84cd9e8d01adb47f046b1ee449fc76a0c32dc4e2/examples/custom_diffusion/train_custom_diffusion.py#L1048) важно обновлять только эмбеддинги для того концепта, который вы пытаетесь выучить. Это означает, что градиенты всех остальных эмбеддингов токенов должны быть установлены в ноль:

```py
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)
```

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

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

В этом руководстве вы скачаете и используете эти примеры [изображений кошек](https://www.cs.cmu.edu/~custom-diffusion/assets/data.zip). При желании вы также можете создать и использовать собственный набор данных (см. руководство [Создание набора данных для обучения](create_dataset)).

Установите переменную окружения `MODEL_NAME` в идентификатор модели на Hub или путь к локальной модели, `INSTANCE_DIR` — в путь, куда вы только что скачали изображения с котами, и `OUTPUT_DIR` — в место, где вы хотите сохранить модель. В качестве специального слова, привязывающего вновь обученные эмбеддинги, используйте `<new1>`. Скрипт создаёт и сохраняет контрольные точки модели и файл pytorch_custom_diffusion_weights.bin в вашем репозитории.
Для отслеживания хода обучения с помощью Weights and Biases добавьте параметр `--report_to=wandb` к команде обучения и укажите промпт для валидации с помощью `--validation_prompt`. Это полезно для отладки и сохранения промежуточных результатов.
> [!TIP]
> Если вы обучаете модель на изображениях человеческих лиц, команда Custom Diffusion обнаружила, что следующие параметры работают хорошо:
>
> - `--learning_rate=5e-6` - `--max_train_steps` может быть в диапазоне от 1000 до 2000 - `--freeze_model=crossattn` - используйте не менее 15-20 изображений для обучения
```bash
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"
export INSTANCE_DIR="./data/cat"

accelerate launch train_custom_diffusion.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --class_data_dir=./real_reg/samples_cat/ \
  --with_prior_preservation \
  --real_prior \
  --prior_loss_weight=1.0 \
  --class_prompt="cat" \
  --num_class_images=200 \
  --instance_prompt="photo of a <new1> cat"  \
  --resolution=512  \
  --train_batch_size=2  \
  --learning_rate=1e-5  \
  --lr_warmup_steps=0 \
  --max_train_steps=250 \
  --scale_lr \
  --hflip  \
  --modifier_token "<new1>" \
  --validation_prompt="<new1> cat sitting in a bucket" \
  --report_to="wandb" \
  --push_to_hub
```
Custom Diffusion также может обучаться нескольким концепциям, если вы предоставите файл [JSON](https://github.com/adobe-research/custom-diffusion/blob/main/assets/concept_list.json) с некоторыми деталями о каждой концепции, которую нужно выучить.
Запустите clip-retrieval, чтобы собрать реальные изображения для использования в качестве регуляризации:
```bash
pip install clip-retrieval
python retrieve.py --class_prompt {} --class_data_dir {} --num_class_images 200
```
Затем вы можете запустить скрипт:
```bash
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"
accelerate launch train_custom_diffusion.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --output_dir=$OUTPUT_DIR \
  --concepts_list=./concept_list.json \
  --with_prior_preservation \
  --real_prior \
  --prior_loss_weight=1.0 \
  --resolution=512  \
  --train_batch_size=2  \
  --learning_rate=1e-5  \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --num_class_images=200 \
  --scale_lr \
  --hflip  \
  --modifier_token "<new1>+<new2>" \
  --push_to_hub
```
Когда обучение завершено, вы можете использовать вашу новую модель Custom Diffusion для инференса.
```py
import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion("path-to-save-model", weight_name="<new1>.bin")

image = pipeline(
    "<new1> cat sitting in a bucket",
    num_inference_steps=100,
    guidance_scale=6.0,
    eta=1.0,
).images[0]
image.save("cat.png")
```
```py
import torch
from huggingface_hub.repocard import RepoCard
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
model_id = "sayakpaul/custom-diffusion-cat-wooden-pot"
pipeline.unet.load_attn_procs(model_id, weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion(model_id, weight_name="<new1>.bin")
pipeline.load_textual_inversion(model_id, weight_name="<new2>.bin")
image = pipeline(
    "the <new1> cat sculpture in the style of a <new2> wooden pot",
    num_inference_steps=100,
    guidance_scale=6.0,
    eta=1.0,
).images[0]
image.save("multi-subject.png")
```

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

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

- Ознакомьтесь с [статьей в блоге о многоконцептуальной кастомизации текстово-изображенийых диффузионных моделей](https://www.cs.cmu.edu/~custom-diffusion/), чтобы узнать больше о результатах экспериментов команды Custom Diffusion.
Xet Storage Details

Size:: 22.1 kB
Xet hash:: 583239d6ee8257e8eada7ee43af81927ec7e3b4984179f2c2f04390ab281594d
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.