Техническое руководство по получению лучших результатов стабильной диффузии.

Это руководство предназначено для тех, кто понимает основы Dreambooth, занимается обучением моделей и хочет добиться лучших результатов на своих моделях. Создавая модели, я разработал несколько лучших практик и идей.

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

Я не буду рассматривать, как работает Dreambooth или где его запустить. Некоторые замечательные ресурсы включают страницу хакатона мечты, бумагу мечты и класс диффузоров с обнимающимся лицом (блок 3). Полный список смотрите в разделе Ресурсы внизу.

Настраивать

Я использовал блокнот huggingface sd_dreambooth_training для тех, кто хочет следовать за мной.

Подсказка экземпляра

Во-первых, нужна подсказка темы. Здесь есть две ключевые переменные: instance_prompt и class_prompt. instance_prompt содержит как ваш уникальный идентификатор, так и идентификатор класса. class_prompt содержит только идентификатор класса. Использование "фотографии" помогает изображениям быть менее абстрактными, особенно при создании изображений классов.

Вот некоторые примеры:

instance_prompt = "photo of jefsnacker person"
prior_preservation_class_prompt = "photo of a person"
instance_prompt = "photo of azzy cat"
prior_preservation_class_prompt = "photo of a cat"
# leva is short for lunar extra vehicular activity (moon landscape)
instance_prompt = "leva landscape"
prior_preservation_class_prompt = "a landscape"
# sks is the original instance prompt used in the dreambooth paper
instance_prompt = "sks dog"
prior_preservation_class_prompt = "a dog"

Это должно не выглядеть так:

instance_prompt = "azzy"
prior_preservation_class_prompt = "cat"

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

Тематические изображения

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

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

Проблема 1. Глаза и рот

Обратите внимание, что на некоторых фотографиях глаза Аззи закрыты, а на других открыты? Представьте себе обученную модель, пытающуюся воссоздать картинку. Он может выбрать глаза, которые наполовину открыты, он может выбрать закрытые, а когда модель переобучится, он может попытаться сделать и то, и другое одновременно. Все это может плохо выглядеть. Это понятие также относится к макияжу и мимике. Старайтесь, чтобы эти вещи были одинаковыми на всех тренировочных фотографиях.

Вердикт: держите лица объектов одинаковыми (глаза, рот, макияж и т. д.). Это нормально, если 5% тренировочных изображений отличаются, но далеко не половина.

Проблема 2. Предыстория

Обратите внимание, что на первых двух изображениях на заднем плане есть пушистое серо-оранжевое существо? По мере того, как модель становится более подходящей, на выходных изображениях могут начать появляться повторяющиеся фоновые узоры, особенно если они выглядят как часть объекта (например, его мех).

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

Проблема 3. Важные функции

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

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

Проблема 4:ТемаРазмер

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

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

Новый тренировочный набор. Имея это в виду, вот улучшенный набор тренировочных изображений.

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

Обрезка.Модель рассчитывает на определенный размер изображения (обычно 512 x 512, но обязательно проверьте!). Если они не такого размера, они будут автоматически обрезаны и изменены. У вас будет гораздо больше контроля, если изображения будут обрезаны перед тренировкой.

Количество фотографий. Вам потребуется более 5, но я не вижу особых улучшений выше 20. У некоторых из моих лучших моделей лиц было около 12, и я видел отличные фотографии. онлайн из 8 и 50 изображений.

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

azzy_01.jpg
azzy_02.jpg
azzy_03.jpg
...

Изображения класса

Большинство сценариев обучения включают раздел для создания изображений класса. Их можно предварительно сгенерировать или загрузить (что экономит МНОГО времени). Я открыл git-репозиторий изображений классов Dreambooth, в котором есть большие наборы изображений классов. При использовании/генерировании изображений классов следует помнить о важных факторах:

  • Используйте достаточное количество изображений, чтобы ваша модель не соответствовала изображениям класса (она может переоснащаться, потому что обучение чередует эти изображения с изображениями предмета). Используйте как минимум вдвое больше изображений классов. Чем больше, тем лучше. Я обычно использую 1500.
  • Дважды проверьте разрешение! Образы классов должны быть такими же, как у модели, которая подвергается точной настройке. Может возникнуть та же проблема, что и с тренировочными изображениями.

Обучение

Время тренироваться, волнительно торопитесь и ждите!

Тренировочные аргументы

Я выложил сюда свою структуру тренировочных аргументов с направляющими комментариями.

Наиболее важные аргументы, на которые следует обратить внимание, это количество_фотографий, learning_rate, max_train_steps и save_steps. Все это коррелирует друг с другом. Я обсуждаю их после структуры args.

args = Namespace(
    # The model to be fine-tuned. Something like stable-diffusion-XXX
    pretrained_model_name_or_path=pretrained_model_name_or_path,

    # Make sure your training images match the model output images
    # Most of the time it's 512, but some models output 768 (sd-2.1)
    resolution=vae.sample_size,

    # What to do if training images don't match 'resolution'
    # shouldn't happen if you followed this guide ;)
    center_crop=True,

    # Very important to be True, 
    # otherwise the text encoder won't learn the new instance prompt word.
    train_text_encoder=True,

    # Where to save output
    instance_data_dir=save_path,

    # What are your training images of?
    instance_prompt=instance_prompt,

    # I've found 1e-06 to be the best. Any higher and the model learns too fast.
    # (at least for complicated things like faces, simple things could be higher)
    learning_rate=1e-06,

    # Approx Rule: num_photos * 100 (for lr=2e-6), photos * 200 (for lr=1e-6)
    # Possibly up to photos * 300 (for lr=1e-6) [faces look better when overfit]
    # NOTE: Discussed below
    max_train_steps=3000,

    # How frequently to save. Takes longer, but more options for best model.
    save_steps=200,

    # How many images to train at a time. Set to 1 if using prior preservation.
    train_batch_size=1, 
    gradient_accumulation_steps=1,

    # Training configs that I didn't touch.
    max_grad_norm=1.0,
    mixed_precision="fp16", # Set to "fp16" for mixed-precision training.

    # Set this to True to lower the memory usage.
    gradient_checkpointing=True, 

    # True: Reduce RAM and might degrade performance.
    use_8bit_adam=False,

    # For reproducibility.
    seed=3434554,

    # Set to True to use class images for training.
    with_prior_preservation=prior_preservation, 

    # I've always used 1.0 for this, never tried anything else.
    prior_loss_weight=prior_loss_weight,

    # 2 and 4 both give stable results, didn't try anything else.
    sample_batch_size=4,

    # Path to class images.
    class_data_dir=prior_preservation_class_folder, 

    # What the class images are of.
    class_prompt=prior_preservation_class_prompt, 

    # Make sure there are enough class images!
    num_class_images=num_class_images, 

    # Flat lr curve because we're fine-tuning.
    lr_scheduler="constant",

    # During warmup a higher LR can be used. 
    # Because we're fine-tuning we don't need this.
    lr_warmup_steps=0,

    # Where to save checkpoints and the model.
    output_dir="dreambooth-concept",
)

Планировщик

Мне удалось успешно использовать DDPMScheduler во время обучения и DPMSolverMultistepSchedulerдля вывода. Интересно, что скрипт использует один для обучения и один для логического вывода. Я немного смущен тем, почему, поэтому, если кто-нибудь знает, пожалуйста, направьте меня в правильном направлении.

# Inference
pipe = StableDiffusionPipeline.from_pretrained(
        model,
        scheduler = DPMSolverMultistepScheduler.from_pretrained(model, subfolder="scheduler"),
        torch_dtype=torch.float16,
    ).to("cuda")

Это сообщение в блоге Huggingface подробно описывает различные планировщики и их влияние.

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

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

# Gradient descent update rule:
Weights = Weights - learning_rate * grad

# Fixing grad leads to:
Weights = Weights - learning_rate * C

# Therefore:
Weights - 2*(learning_rate_1e-6 * C) == Weights - (learning_rate_2e-6 * C)

# Basically, 2 steps at 1e-6 is similar to 1 step at 2e-6.

Короче говоря, более низкая learning_rate означает, что требуется больше шагов.

Интернет-ресурсы сошлись во мнении, что 1e-6 — лучший вариант для лиц. Более высокие скорости обучения могут использоваться для менее подробных предметов. FollowFox предлагает отличное сравнение скорости обучения и количества шагов.

Количество фотографий

Следующая переменная, о которой следует подумать, — это num_photos. Если предположить, что все тренировочные фотографии уникальны, то количество раз, которое модель «видит» каждую фотографию во время тренировки, — это то, что нас волнует, тогда мы можем сказать следующее:

# num_times_seen   - number of times to see each photo
# num_photos       - number of photos in the train
# train_batch_size - number of photos used at each training step
steps_needed = num_photos * num_times_seen / train_batch_size

Короче говоря, с более уникальными фотографиями требуется больше шагов обучения.

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

Собираем все вместе

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

# For learning faces
step_guess = num_photos * 300 / (learning rate * 1e6)

# Examples
step_guess = (12 photos * 300) / (1e-6 * 1e6) = 3600
step_guess = (6 photos * 300) / (2e-6 * 1e6) = 900

Если кто-то знает лучшие уравнения для ландшафтов, еды или чего-то еще, поделитесь!

Единственная проблема заключается в том, что аргументы принимают не step_guess, а max_train_steps и save_steps. Теперь мы хотим выбрать эти переменные, чтобы иметь несколько контрольных точек на выбор, которые близки к этому идеальному количеству шагов.

Вот несколько примеров того, что я могу выбрать:

# step_guess = 3600 (lr=1e-6, num_photos=12)
# save checkpoints at 600, 1200, 1800, 2400, 3000, 3600, 4200, 4800 
max_train_steps = 4800
save_steps = 600


# step_guess = 900 (lr=2e-6, num_photos=6)
# save checkpoints at 200, 400, 600, 800, 1000, 1200, 1400
max_train_steps = 1400
save_steps = 200

Я обычно превышаю step_guess на 30–40 %, поэтому можно увидеть, как модель переходит в состояние переобучения. Хотя создание множества контрольных точек кажется заманчивым, важно отметить, что контрольные точки занимают несколько ГБ каждая и замедляют время обучения.

Посттренинг

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

Выбор контрольной точки

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

prompt = "portrait of azzy cat swimming underwater"
num_samples = 6
guidance_scale = 8
num_inference_steps = 50

# Paths to all checkpoints.
# Note: Google Colab has a hard time displaying more than 4 rows at a time.
model_list = [
    # 'dreambooth-concept/checkpoint-200',
    'dreambooth-concept/checkpoint-400',
    # 'dreambooth-concept/checkpoint-600',
    'dreambooth-concept/checkpoint-800',
    # 'dreambooth-concept/checkpoint-1000',
    'dreambooth-concept/checkpoint-1200',
    # 'dreambooth-concept/checkpoint-1400',
    'dreambooth-concept/checkpoint-1600',
]

# output images
all_images = []
for model in model_list:
    # Setup pipeline for checkpoint
    pipe = StableDiffusionPipeline.from_pretrained(
        model,
        scheduler = DPMSolverMultistepScheduler.from_pretrained(model, subfolder="scheduler"),
        torch_dtype=torch.float16,
    ).to("cuda")

    # Set the seed to compare checkpoints.
    generator = torch.Generator(device="cuda").manual_seed(42)

    # Generate images & add to output
    images = pipe(
        prompt, 
        num_images_per_prompt=num_samples, 
        num_inference_steps=num_inference_steps, 
        guidance_scale=guidance_scale,
        generator=generator,
    ).images
    all_images.extend(images)

# Display all in a grid
grid = image_grid(all_images, len(model_list), num_samples)
grid

Я возьму пример с моим котом Азриэлем. Приведенный выше скрипт создаст сетку, подобную следующей (используя другую подсказку и контрольные точки):

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

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

Контрольная точка где-то посередине лучше всего. В данном случае ряд 2.

Испытание множества подсказок жизненно важно для выбора лучшей контрольной точки!

Подсказка

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

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

# Training prompts
instance_prompt = "photo of azzy cat"
class_prompt = "photo of a cat"

# Will strongly depict azzy by using "azzy cat"
prompt = "azzy cat wearing a tuxedo at a wedding"

## Trick 1: inclusion/exclusion of class prompt
# Will portray azzy in a more general way, not as strong
prompt = "azzy wearing a tuxedo at a wedding"

## Trick 2: placement of instance key word
# Weak portrayal. Use if "azzy" is stronger than the rest of the prompt 
prompt = "at a wedding with azzy cat in a tuxedo" 

## Trick 3: weighting
# This is built into the model generation, and can accomdate negative numbers.
prompt = "(azzy cat:0.4) wearing a tuxedo at a wedding"

Чтобы получить наилучшие результаты, обычно необходимо опробовать множество вариантов.

Заканчивать

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

Ресурсы

Это набор ресурсов, которые я нашел полезными во время обучения. Этот список далеко не полный.

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

Руководства по обучению:

Сайты для вдохновения в искусстве искусственного интеллекта: