Hugging Face's logo

Diffusers documentation

DreamBooth

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

DreamBooth

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

Чтобы загрузить обученный контрольный пункт для вывода, см. Загрузка модели DreamBooth для вывода.

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

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

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

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

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

cd examples/dreambooth
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, чтобы узнать, как создать набор данных, совместимый со скриптом обучения.

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

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

DreamBooth is very sensitive to training hyperparameters, and it is easy to overfit. Read the Training Stable Diffusion with Dreambooth using 🧨 Diffusers blog post for recommended settings for different subjects to help you choose the appropriate hyperparameters.

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

Например, для обучения в формате bf16:

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

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

  • --pretrained_model_name_or_path: название модели на Hub или локальный путь к предобученной модели
  • --instance_data_dir: путь к папке, содержащей обучающий набор данных (примеры изображений)
  • --instance_prompt: текстовый промпт, содержащий специальное слово для примерных изображений
  • --train_text_encoder: нужно ли также обучать текстовый энкодер
  • --output_dir: место для сохранения обученной модели
  • --push_to_hub: нужно ли отправить обученную модель в Hub
  • --checkpointing_steps: частота сохранения контрольных точек в процессе обучения модели; это полезно, если обучение по какой-либо причине прерывается, вы можете продолжить обучение с этой контрольной точки, добавив --resume_from_checkpoint к команде обучения

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

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

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

accelerate launch train_dreambooth.py \
  --snr_gamma=5.0

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

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

  • --with_prior_preservation: использовать ли потерю сохранения априорного распределения
  • --prior_loss_weight: управляет влиянием функции потерь сохранения априорного распределения на модель
  • --class_data_dir: путь к папке, содержащей сгенерированные образцы изображений классов
  • --class_prompt: текстовый промпт, описывающий класс генерируемых изображений
accelerate launch train_dreambooth.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="path/to/class/images" \
  --class_prompt="text prompt describing class"

Обучение текстового кодировщика

Для повышения качества генерируемых результатов можно также обучать текстовый энкодер в дополнение к UNet. Для этого требуется дополнительное количество памяти, а GPU должна иметь не менее 24 ГБ vRAM. Если у вас есть подходящее оборудование, обучение текстового энкодера даёт лучшие результаты, особенно при генерации изображений лиц. Включите этот параметр, выполнив следующие действия:

accelerate launch train_dreambooth.py \
  --train_text_encoder

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

DreamBooth поставляется со своими собственными классами наборов данных:

  • DreamBoothDataset: выполняет предобработку изображений и класс-изображений, а также токенизацию промптов для обучения
  • PromptDataset: генерирует эмбеддинги промптов для создания изображений классов

Если вы включили prior preservation loss, класс-изображения генерируются здесь:

sample_dataset = PromptDataset(args.class_prompt, num_new_images)
sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)

sample_dataloader = accelerator.prepare(sample_dataloader)
pipeline.to(accelerator.device)

for example in tqdm(
    sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
):
    images = pipeline(example["prompt"]).images

Далее идёт функция main(), которая отвечает за подготовку набора данных для обучения и сам цикл обучения. Скрипт загружает токенизатор, планировщик и модели:

# Load the tokenizer
if args.tokenizer_name:
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, revision=args.revision, use_fast=False)
elif args.pretrained_model_name_or_path:
    tokenizer = AutoTokenizer.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="tokenizer",
        revision=args.revision,
        use_fast=False,
    )

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = text_encoder_cls.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)

if model_has_vae(args):
    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision
    )
else:
    vae = None

unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

Затем пришло время создать обучающий набор данных и DataLoader из DreamBoothDataset:

train_dataset = DreamBoothDataset(
    instance_data_root=args.instance_data_dir,
    instance_prompt=args.instance_prompt,
    class_data_root=args.class_data_dir if args.with_prior_preservation else None,
    class_prompt=args.class_prompt,
    class_num=args.num_class_images,
    tokenizer=tokenizer,
    size=args.resolution,
    center_crop=args.center_crop,
    encoder_hidden_states=pre_computed_encoder_hidden_states,
    class_prompt_encoder_hidden_states=pre_computed_class_prompt_encoder_hidden_states,
    tokenizer_max_length=args.tokenizer_max_length,
)

train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=args.train_batch_size,
    shuffle=True,
    collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
    num_workers=args.dataloader_num_workers,
)

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

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

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

Теперь вы готовы запустить скрипт обучения! 🚀

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

from huggingface_hub import snapshot_download

local_dir = "./dog"
snapshot_download(
    "diffusers/dog-example",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

Установите переменную окружения MODEL_NAME в значение идентификатора модели на Hub или пути к локальной модели, INSTANCE_DIR — в путь, куда вы только что скачали изображения собак, и OUTPUT_DIR — в место, где вы хотите сохранить модель. Вы будете использовать sks как специальное слово для привязки обучения.

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

--validation_prompt="a photo of a sks dog"
--num_validation_images=4
--validation_steps=100

Ещё одно замечание перед запуском скрипта! В зависимости от вашей видеокарты может потребоваться включить определённые оптимизации для обучения DreamBooth.

16 ГБ
12 ГБ
8 ГБ

На GPU с 16 ГБ памяти можно использовать 8-битный оптимизатор bitsandbytes и gradient checkpointing, чтобы помочь в обучении модели DreamBooth. Установите bitsandbytes:

pip install bitsandbytes

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

accelerate launch train_dreambooth.py \
  --gradient_checkpointing \
  --use_8bit_adam \
export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export INSTANCE_DIR="./dog"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=400 \
  --push_to_hub

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

Не можете дождаться, чтобы попробовать вашу модель для инференса до завершения обучения? 🤭 Убедитесь, что у вас установлена последняя версия 🤗 Accelerate.

from diffusers import DiffusionPipeline, UNet2DConditionModel
from transformers import CLIPTextModel
import torch

unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet")

# if you have trained with `--args.train_text_encoder` make sure to also load the text encoder
text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder")

pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16,
).to("cuda")

image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")
from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path_to_saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

LoRA

LoRA — это метод обучения, позволяющий значительно сократить количество обучаемых параметров. В результате обучение становится быстрее, а хранение полученных весов упрощается, так как их размер существенно меньше (~100 МБ). Для обучения с использованием LoRA используйте скрипт train_dreambooth_lora.py.

Скрипт обучения LoRA подробно описан в руководстве Обучение LoRA.

Stable Diffusion XL

Stable Diffusion XL (SDXL) — это мощная модель text-to-image, генерирующая изображения высокого разрешения, которая добавляет второй text-encoder в свою архитектуру. Используйте скрипт train_dreambooth_lora_sdxl.py для обучения модели SDXL с LoRA.

Скрипт обучения SDXL подробно описан в руководстве обучение SDXL.

DeepFloyd IF

DeepFloyd IF — это каскадная пиксельная диффузионная модель с тремя этапами. Первый этап генерирует базовое изображение, а второй и третий этапы последовательно увеличивают разрешение базового изображения до высокого разрешения 1024x1024. Используйте скрипты train_dreambooth_lora.py или train_dreambooth.py для обучения модели DeepFloyd IF с использованием LoRA или полной модели.

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

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", use_safetensors=True)

pipe.load_lora_weights("<lora weights path>")

# Update scheduler config to fixed variance schedule
pipe.scheduler = pipe.scheduler.__class__.from_config(pipe.scheduler.config, variance_type="fixed_small")

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

from huggingface_hub import snapshot_download

local_dir = "./dog_downsized"
snapshot_download(
    "diffusers/dog-example-downsized",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

Приведённые ниже примеры кода кратко демонстрируют, как обучить модель DeepFloyd IF с использованием комбинации DreamBooth и LoRA. Обратите внимание на некоторые важные параметры:

  • --resolution=64, требуется гораздо меньшее разрешение, так как DeepFloyd IF — это пиксельная диффузионная модель, и для работы с несжатыми пикселями входные изображения должны быть меньшего размера
  • --pre_compute_text_embeddings, вычисляет текстовые эмбеддинги заранее, чтобы сэкономить память, так как T5Model может занимать значительный объём памяти
  • --tokenizer_max_length=77, вы можете использовать более длинную длину текста по умолчанию с T5 в качестве текстового энкодера, но процедура кодирования по умолчанию в модели использует более короткую длину текста
  • --text_encoder_use_attention_mask, чтобы передать маску внимания в текстовый энкодер
Этап 1 LoRA DreamBooth
Этап 2 LoRA DreamBooth
Этап 1 DreamBooth
Этап 2 DreamBooth

Обучение первого этапа DeepFloyd IF с использованием LoRA и DreamBooth требует ~28 ГБ памяти.

export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth_dog_lora"

accelerate launch train_dreambooth_lora.py \
  --report_to wandb \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a sks dog" \
  --resolution=64 \
  --train_batch_size=4 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --scale_lr \
  --max_train_steps=1200 \
  --validation_prompt="a sks dog" \
  --validation_epochs=25 \
  --checkpointing_steps=100 \
  --pre_compute_text_embeddings \
  --tokenizer_max_length=77 \
  --text_encoder_use_attention_mask

Советы по обучению

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

  • LoRA достаточно для обучения модели первого этапа, поскольку низкое разрешение модели затрудняет представление мелких деталей в любом случае.
  • Для общих или простых объектов не обязательно выполнять дообучение апскейлера. Убедитесь, что промпт, передаваемый апскейлеру, скорректирован так, чтобы новый токен был удалён из промпта экземпляра. Например, если ваш промпт первого этапа — «a sks dog», то промпт второго этапа должен быть «a dog».
  • Для более тонких деталей, таких как лица, полная дообучение апскейлера второго этапа даёт лучший результат, чем дообучение модели второго этапа с использованием LoRA. Также помогает использовать более низкие скорости обучения при больших размерах батчей.
  • Для обучения модели второго этапа следует использовать более низкие значения коэффициента обучения.
  • DDPMScheduler работает лучше, чем DPMSolver, используемый в скриптах обучения.

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

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

Update on GitHub

←Textual Inversion LoRA→