Textual Inversion

Textual Inversion — это метод обучения для персонализации моделей генерации изображений на основе нескольких примеров изображений, которые модель должна освоить. Метод работает за счёт обучения и обновления текстовых эмбеддингов (новые эмбеддинги привязываются к специальному слову, которое необходимо использовать в промпте) так, чтобы они соответствовали предоставленным вами примерам изображений.

Для инференса с обученными эмбеддингами см. Инференс textual inversion.

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

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

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

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

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

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

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

accelerate launch textual_inversion.py \
  --gradient_accumulation_steps=4

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

• --pretrained_model_name_or_path: название модели на Hub или локальный путь к предобученной модели
• --train_data_dir: путь к папке, содержащей обучающий набор данных (примеры изображений)
• --output_dir: место для сохранения обученной модели
• --push_to_hub: нужно ли отправить обученную модель в Hub
• --checkpointing_steps: частота сохранения контрольных точек в процессе обучения модели; это полезно, если обучение по какой-либо причине прерывается, вы можете продолжить обучение с этой контрольной точки, добавив --resume_from_checkpoint к команде обучения
• --num_vectors: количество векторов, используемых для обучения эмбеддингов; увеличение этого параметра помогает модели обучаться лучше, но приводит к росту затрат на обучение
• --placeholder_token: специальное слово, к которому привязываются обученные эмбеддинги (это слово должно присутствовать в вашем промпте при инференсе)
• --initializer_token: одно слово, которое примерно описывает объект или стиль, на котором вы хотите обучить модель
• --learnable_property: указывает, обучаете ли вы модель новому «стилю» (например, стилю картин Ван Гога) или «объекту» (например, вашей собаке)

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

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

Далее в функции main() вы найдёте код предобработки набора данных и цикл обучения.

Скрипт начинается с загрузки токенизатора, планировщика и модели:

# Load tokenizer
if args.tokenizer_name:
    tokenizer = CLIPTokenizer.from_pretrained(args.tokenizer_name)
elif args.pretrained_model_name_or_path:
    tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer")

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = CLIPTextModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

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

Затем скрипт создает набор данных из TextualInversionDataset:

train_dataset = TextualInversionDataset(
    data_root=args.train_data_dir,
    tokenizer=tokenizer,
    size=args.resolution,
    placeholder_token=(" ".join(tokenizer.convert_ids_to_tokens(placeholder_token_ids))),
    repeats=args.repeats,
    learnable_property=args.learnable_property,
    center_crop=args.center_crop,
    set="train",
)
train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.dataloader_num_workers
)

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

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

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

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

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

from huggingface_hub import snapshot_download

local_dir = "./cat"
snapshot_download(
    "diffusers/cat_toy_example", local_dir=local_dir, repo_type="dataset", ignore_patterns=".gitattributes"
)

Установите переменную окружения MODEL_NAME в идентификатор модели на Hub или путь к локальной модели, а DATA_DIR — в путь, куда вы только что скачали изображения кошек. Скрипт создаёт и сохраняет в ваш репозиторий следующие файлы:

• learned_embeds.bin: обученные векторы эмбеддингов, соответствующие вашим примерам изображений
• token_identifier.txt: специальный токен-заполнитель
• type_of_concept.txt: тип концепции, на которой выполняется обучение (либо «object», либо «style»)

Полный цикл обучения занимает около 1 часа на одном GPU V100.

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

--validation_prompt="A <cat-toy> train"
--num_validation_images=4
--validation_steps=100

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export DATA_DIR="./cat"

accelerate launch textual_inversion.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATA_DIR \
  --learnable_property="object" \
  --placeholder_token="<cat-toy>" \
  --initializer_token="toy" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=3000 \
  --learning_rate=5.0e-04 \
  --scale_lr \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --output_dir="textual_inversion_cat" \
  --push_to_hub

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

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.load_textual_inversion("sd-concepts-library/cat-toy")
image = pipeline("A <cat-toy> train", num_inference_steps=50).images[0]
image.save("cat-train.png")

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

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

• Узнайте, как загружать эмбеддинги Textual Inversion и использовать их в качестве отрицательных эмбеддингов.