Stable Diffusion XL

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

Stable Diffusion XL (SDXL) — это более крупная и мощная итерация модели Stable Diffusion, способная генерировать изображения с более высоким разрешением.

UNet в SDXL в 3 раза больше, а в архитектуру модели добавлен второй текстовый энкодер. В зависимости от доступного вам оборудования это может требовать значительных вычислительных ресурсов, и модель может не запуститься на потребительском GPU, таком как Tesla T4. Чтобы уместить эту более крупную модель в память и ускорить обучение, попробуйте включить gradient_checkpointing, mixed_precision и gradient_accumulation_steps. Вы можете ещё больше снизить потребление памяти, включив энергоэффективное внимание с помощью xFormers и используя 8-битный оптимизатор bitsandbytes.

This guide will explore the train_text_to_image_sdxl.py training script to help you become more familiar with it, and how you can adapt it for your own use-case.

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

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

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

cd examples/text_to_image
pip install -r requirements_sdxl.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(). Эта функция задаёт значения по умолчанию для каждого параметра, например, размер обучающей выборки и скорость обучения, но при необходимости вы можете указать свои значения в команде запуска обучения.

Например, чтобы ускорить обучение с использованием смешанной точности в формате bf16, добавьте параметр --mixed_precision к команде обучения:

accelerate launch train_text_to_image_sdxl.py \
  --mixed_precision="bf16"

Большинство параметров совпадают с параметрами из руководства по обучению Text-to-image, поэтому в данном руководстве вы сосредоточитесь на параметрах, относящихся к обучению SDXL.

• --pretrained_vae_model_name_or_path: путь к предобученной VAE; известно, что VAE в SDXL страдает от численной нестабильности, поэтому этот параметр позволяет указать более подходящую VAE
• --proportion_empty_prompts: доля текстовых подсказок для изображений, которые заменяются пустыми строками
• --timestep_bias_strategy: где (на более ранних или более поздних шагах) применять смещение, что может побуждать модель обучаться деталям с низкой или высокой частотой
• --timestep_bias_multiplier: вес смещения, применяемого к шагу
• --timestep_bias_begin: шаг времени, с которого начинает применяться смещение
• --timestep_bias_end: шаг времени, на котором прекращается применение смещения
• --timestep_bias_portion: доля шагов, к которым применяется смещение

Взвешивание Min-SNR

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

Добавьте параметр --snr_gamma и установите для него рекомендуемое значение 5.0:

accelerate launch train_text_to_image_sdxl.py \
  --snr_gamma=5.0

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

Скрипт обучения также похож на руководство по обучению Text-to-image, но был изменён для поддержки обучения SDXL. В этом руководстве мы сосредоточимся на коде, который уникален для скрипта обучения SDXL.

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

В функции main(), помимо загрузки токенизатора, скрипт загружает второй токенизатор и текстовый энкодер, поскольку архитектура SDXL использует по два каждого:

tokenizer_one = AutoTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision, use_fast=False
)
tokenizer_two = AutoTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer_2", revision=args.revision, use_fast=False
)

text_encoder_cls_one = import_model_class_from_model_name_or_path(
    args.pretrained_model_name_or_path, args.revision
)
text_encoder_cls_two = import_model_class_from_model_name_or_path(
    args.pretrained_model_name_or_path, args.revision, subfolder="text_encoder_2"
)

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

text_encoders = [text_encoder_one, text_encoder_two]
tokenizers = [tokenizer_one, tokenizer_two]
compute_embeddings_fn = functools.partial(
    encode_prompt,
    text_encoders=text_encoders,
    tokenizers=tokenizers,
    proportion_empty_prompts=args.proportion_empty_prompts,
    caption_column=args.caption_column,
)

train_dataset = train_dataset.map(compute_embeddings_fn, batched=True, new_fingerprint=new_fingerprint)
train_dataset = train_dataset.map(
    compute_vae_encodings_fn,
    batched=True,
    batch_size=args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps,
    new_fingerprint=new_fingerprint_for_vae,
)

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

del text_encoders, tokenizers, vae
gc.collect()
torch.cuda.empty_cache()

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

weights = generate_timestep_weights(args, noise_scheduler.config.num_train_timesteps).to(
        model_input.device
    )
    timesteps = torch.multinomial(weights, bsz, replacement=True).long()

noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)

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

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

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

Обучимся на датасете Naruto BLIP captions, чтобы генерировать собственных персонажей Наруто. Установите переменные окружения MODEL_NAME и DATASET_NAME в значения модели и датасета (из Hub или локальный путь). Также с помощью VAE_NAME укажите VAE, отличный от SDXL VAE (из Hub или локальный путь), чтобы избежать численных нестабильностей.

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

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fix"
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch train_text_to_image_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --pretrained_vae_model_name_or_path=$VAE_NAME \
  --dataset_name=$DATASET_NAME \
  --enable_xformers_memory_efficient_attention \
  --resolution=512 \
  --center_crop \
  --random_flip \
  --proportion_empty_prompts=0.2 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=10000 \
  --use_8bit_adam \
  --learning_rate=1e-06 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --mixed_precision="fp16" \
  --report_to="wandb" \
  --validation_prompt="a cute Sundar Pichai creature" \
  --validation_epochs 5 \
  --checkpointing_steps=5000 \
  --output_dir="sdxl-naruto-model" \
  --push_to_hub

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

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path/to/your/model", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A naruto with green eyes and red legs."
image = pipeline(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("naruto.png")

from diffusers import DiffusionPipeline
import torch
import torch_xla.core.xla_model as xm

device = xm.xla_device()
pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to(device)

prompt = "A naruto with green eyes and red legs."
start = time()
image = pipeline(prompt, num_inference_steps=inference_steps).images[0]
print(f'Compilation time is {time()-start} sec')
image.save("naruto.png")

start = time()
image = pipeline(prompt, num_inference_steps=inference_steps).images[0]
print(f'Inference time is {time()-start} sec after compilation')

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

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

• Ознакомьтесь с руководством по Stable Diffusion XL, чтобы узнать, как использовать его для различных задач (генерация изображений по тексту, преобразование изображения в изображение, inpainting), как применять его модель-рефинер, а также о различных типах микро-условий.
• Ознакомьтесь с руководствами по обучению DreamBooth и LoRA, чтобы узнать, как обучить персонализированную модель SDXL всего на нескольких примерах изображений. Эти два метода обучения можно даже комбинировать!