Персонализированные сгенерированные изображения с пользовательскими стилями или объектами

Ранее я публиковал статью о тонкой настройке стабильной диффузии с помощью текстовой инверсии. В этом руководстве основное внимание уделяется тонкой настройке стабильной диффузии с использованием другого метода, называемого Dreambooth. В отличие от метода текстовой инверсии, который обучает только встраиванию без изменения базовой модели, Dreambooth точно настраивает всю модель преобразования текста в изображение, чтобы научиться связывать уникальный идентификатор с определенной концепцией (объектом или стилем). В результате сгенерированные изображения более персонализированы для объекта или стиля по сравнению с текстовой инверсией.

Этот туториал основан на разветвленной версии реализации Dreambooth от HuggingFace. Первоначальная реализация требует от 16 ГБ до 24 ГБ для точной настройки модели. Сопровождающий ShivamShirao оптимизировал код, чтобы уменьшить использование видеопамяти до 16 ГБ. В зависимости от ваших потребностей и настроек вы можете настроить модель с графическим процессором от 10 ГБ до 16 ГБ. Я лично протестировал обучение на графическом процессоре Tesla T4.

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

Давайте перейдем к следующему разделу, чтобы настроить все необходимые модули.

Настраивать

Перед продолжением установки рекомендуется создать новую виртуальную среду.

Пакеты Python

В рабочем каталоге создайте новый файл с именем requirements.txt со следующим кодом:

accelerate==0.12.0
torchvision
transformers>=4.21.0
ftfy
tensorboard
modelcards

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

pip install git+https://github.com/ShivamShrirao/diffusers.git
pip install -r requirements.txt

ПРИМЕЧАНИЕ. Вам необходимо установить diffusers, используя приведенный выше URL-адрес, а не устанавливать его напрямую с pypi.

пакет bitsandbytes

Существует дополнительный пакет под названием bitsandbytes, который может еще больше сократить использование видеопамяти. Однако он поддерживает только версии CUDA 10.2–11.7, и ваша машина должна соответствовать следующим требованиям:

  • LLM.int8(): NVIDIA Turing (RTX 20xx; T4) или графический процессор Ampere (RTX 30xx; A4-A100); (графический процессор 2018 года или старше).
  • 8-bit optimizers and quantization: графический процессор NVIDIA Maxwell или новее (›=GTX 9XX).

Вы можете установить его следующим образом:

pip install bitsandbytes

пакет xformers

Для тех, у кого GPU меньше 24 ГБ, вам необходимо установить пакет xformers, чтобы еще больше сократить использование VRAM. На момент написания этой статьи установка xformers была не такой уж простой из-за отсутствия поддержки колес со стороны разработчиков.

Вы можете установить пакет следующим образом:

pip install xformers

Если вы столкнулись с ошибкой в ​​приведенной выше команде, выполните следующую команду, чтобы собрать пакет непосредственно из репозитория:

pip install git+https://github.com/facebookresearch/xformers.git@main#egg=xformers

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

Ускорить настройку

Следующим шагом является инициализация среды Accelerate. Выполните следующую команду:

accelerate config

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

In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS): 0
Do you want to run your training on CPU only (even if a GPU is available)? [yes/NO]:no
Do you want to use DeepSpeed? [yes/NO]: no
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: fp16

Модели HuggingFace

Вы можете пропустить этот раздел, если у вас уже есть модель diffusers для стабильной диффузии (v1.4/v1.5). Для тех, кто

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

Зарегистрируйте новую учетную запись в HuggingFace и примите лицензию модели перед загрузкой или использованием гирь.

Как только вы закончите с этим, обратитесь к этому разделу документации, чтобы включить токен доступа.

Выполните следующую команду и передайте свой токен для аутентификации:

huggingface-cli login

Он загрузит веса в папку кеша во время начального запуска.

Наборы данных

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

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

В зависимости от ваших вариантов использования вы можете использовать следующие рекомендации:

Объект

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

  • угол камеры
  • поза
  • реквизит (одежда, прическа и т.д.)
  • фон (снято в разных местах)

Количество обучающих изображений должно быть от 5 до 20. Возможно, вам придется обрезать изображения, чтобы сфокусироваться только на объекте.

Стиль

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

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

Тренировочные изображения

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

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

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

Обучение

Перейдите в следующий репозиторий Github и загрузите файл train_dreambooth.py в свой рабочий каталог.

Тренировочные команды

Ниже приведен список общих параметров, которые следует изменить в зависимости от ваших вариантов использования:

  • pretrained_model_name_or_path — Путь к предварительно обученной модели или идентификатору модели из Huggingface.co/models
  • pretrained_vae_name_or_path — Путь к предварительно обученному vae или идентификатору vae из Huggingface.co/models. Вы можете точно настроить модель с vae или без него.
  • instance_data_dir — Папка, содержащая обучающие данные экземпляров изображений.
  • class_data_dir — Папка, содержащая обучающие данные изображений классов.
  • instance_prompt — Подсказка с идентификатором, указывающим экземпляр
  • class_prompt — Приглашение указать изображения того же класса, что и предоставленные изображения экземпляра.
  • num_class_images — Изображения минимального класса для предварительной потери сохранения
  • output_dir — Выходной каталог, в который будут записаны предсказания модели и контрольные точки.
  • max_train_steps — Общее количество шагов обучения для выполнения. Рекомендуется установить его на N * 100, где N представляет количество изображений экземпляров.
  • learning_rate — Начальная скорость обучения (после потенциального периода прогрева) для использования
  • lr_scheduler— Тип планировщика для использования. Выберите между [linear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmup]
  • lr_warmup_steps — Количество шагов для прогрева в планировщике lr. Используйте max_train_steps / 10 при использовании polynomial или 0 при использовании constant.
  • save_interval — Сохранять веса каждые N шагов. Убедитесь, что у вас достаточно места для его хранения. Каждый вес около 4 Гб.

Вы можете установить пользовательский vae, используя следующий флаг:

--pretrained_vae_name_or_path="stabilityai/sd-vae-ft-mse"

Обучение с потерей предварительного сохранения помогает предотвратить переоснащение. Включите его следующим образом:

--with_prior_preservation --prior_loss_weight=1.0

Кроме того, вы можете настроить text_encoder вместе с unet. Однако это резко увеличит использование VRAM. Установите его со следующим флагом:

--train_text_encoder

Для обучения объекта вы можете использовать следующий пример в качестве справки для instance_prompt и class_prompt. Не стесняйтесь экспериментировать с разными строками в зависимости от ваших вариантов использования.

# Woman
--instance_prompt="photo of zwx woman" \
--class_prompt="photo of a woman" \

# Black man
--instance_prompt="photo of zwx black man" \
--class_prompt="photo of a black man" \

# Dog
--instance_prompt="photo of zwx dog" \
--class_prompt="photo of a dog" \

Вы можете использовать свою собственную строку в качестве уникального идентификатора. В ранней реализации в большинстве примеров в качестве уникального идентификатора используется sks. Однако sks — это известный токен полуавтоматической винтовки. Настоятельно рекомендуется использовать другой уникальный идентификатор, который не является частью токена, используемого в исходных наборах данных Stable Diffusion.

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

# Style 1
--instance_prompt="modern disney style" \
--class_prompt="artwork style" \

# Style 2
--instance_prompt="classic animation style" \
--class_prompt="illustration style" \

Что касается скорости обучения и планировщика, пожалуйста, обратитесь к следующему блогу HuggingFace.

Примеры обучения

Взгляните на следующий пример обучающей команды для обучающих изображений выше (проверено на Tesla T4 с 16 ГБ памяти):

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
  --pretrained_vae_name_or_path="stabilityai/sd-vae-ft-mse" \
  --instance_data_dir="./instance-images/" \
  --class_data_dir="./class-images/" \
  --output_dir="./output-models/" \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --instance_prompt="photo of zwx bear toy" \
  --class_prompt="photo of bear toy" \
  --resolution=512 \
  --train_batch_size=1 \
  --train_text_encoder \
  --mixed_precision="fp16" \
  --use_8bit_adam \
  --gradient_accumulation_steps=1 \
  --gradient_checkpointing \
  --learning_rate=1e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=200 \
  --num_class_images=300 \
  --max_train_steps=2000 \
  --save_interval=500

Измените каталог данных соответственно на основе вашего рабочего каталога.

Если у вас недостаточно памяти, отключите тренировку text_encoder:

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
  --instance_data_dir="./instance-images/" \
  --class_data_dir="./class-images/" \
  --output_dir="./output-models/" \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --instance_prompt="photo of zwx bear toy" \
  --class_prompt="photo of bear toy" \
  --resolution=512 \
  --train_batch_size=1 \
  --mixed_precision="fp16" \
  --use_8bit_adam \
  --gradient_accumulation_steps=1 \
  --gradient_checkpointing \
  --learning_rate=1e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=300 \
  --max_train_steps=2000 \
  --save_interval=500

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

Более того, последний обучающий скрипт принимает новый параметр с именем concept_list. Он представляет собой путь к файлу JSON, содержащему список словарей. Он перезапишет такие параметры, как instance_prompt, class_prompt и т. д. Вы можете использовать это для одновременного обучения нескольких концепций одной модели. Например, для следующего файла concept_list.json:

[
    {
        "instance_prompt":      "photo of zwx dog",
        "class_prompt":         "photo of a dog",
        "instance_data_dir":    "./instance-images/",
        "class_data_dir":       "./class-images/"
    }
]

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

--concepts_list ./concepts_list.json

Просто добавьте новый словарь в список для одновременного обучения другому понятию.

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

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

Результат обучения

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

|- feature_extractor
|  |- preprocessor_config.json
|- scheduler
|  |- scheduler_config.json
|- text_encoder
|  |- config.json
|  |- pytorch_model.bin
|- tokenizer
|  |- merges.txt
|  |- special_tokens_map.json
|  |- tokenizer_config.json
|  |- vocab.json
|- unet
|  |- config.json
|  |- diffusion_pytorch_model.bin
|- vae
|  |- config.json
|  |- diffusion_pytorch_model.bin
|- args.json
|- model_index.json

Вывод

Теперь создайте новый файл Python с именем inference.py в своем рабочем каталоге. Добавьте в него следующий код:

from diffusers import StableDiffusionPipeline, DDIMScheduler
import torch

device = "cuda"
# use DDIM scheduler, you can modify it to use other scheduler
scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=True)

# modify the model path
pipe = StableDiffusionPipeline.from_pretrained(
    f"./output-models/1500/",
    scheduler=scheduler,
    safety_checker=None,
    torch_dtype=torch.float16,
).to(device)

# enable xformers memory attention
pipe.enable_xformers_memory_efficient_attention()

prompt = "photo of zwx bear toy"
negative_prompt = ""
num_samples = 4
guidance_scale = 7.5
num_inference_steps = 50
height = 512
width = 512

images = pipe(
    prompt,
    height=height,
    width=width,
    negative_prompt=negative_prompt,
    num_images_per_prompt=num_samples,
    num_inference_steps=num_inference_steps,
    guidance_scale=guidance_scale
).images

count = 1
for image in images:
    # save image to local directory
    image.save(f"img-{count}.png")
    count += 1

Некоторые пользователи сообщали, что использование xformers для создания изображений приведет к недетерминированным результатам. Это означает, что вы не можете воспроизвести одно и то же изображение с одинаковыми настройками (начальное число, размер изображения и т. д.). Пожалуйста, поэкспериментируйте сами и измените код соответствующим образом в зависимости от ваших вариантов использования.

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

python inference.py

Вот некоторые из примеров вывода (подсказки не включены, так как я сделал ошибку, не отследив их изначально):

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

Заключение

Dreambooth — хороший метод тонкой настройки модели Stable Diffusion с определенной концепцией (объектом или стилем).

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

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

К сожалению, на данный момент у меня недостаточно памяти для тонкой настройки всей модели с использованием пар изображений и подписей.

Спасибо за чтение этого произведения. Желаю отличного дня!

Рекомендации

  1. HuggingFace Diffusers Github
  2. Github Dreambooth ШивамШрирао