Latent Consistency Distillation

Latent Consistency Models (LCMs) способны генерировать изображения высокого качества всего за несколько шагов, что является значительным прорывом, поскольку многим пайплайнам требуется как минимум 25 и более шагов. LCMs создаются путём применения метода дистилляции латентной согласованности к любой модели Stable Diffusion. Этот метод работает за счёт применения одностадийной направленной дистилляции к латентному пространству и использования метода skipping-step для последовательного пропуска временных шагов, что ускоряет процесс дистилляции (подробнее см. разделы 4.1, 4.2 и 4.3 статьи).

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

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

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

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

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

cd examples/consistency_distillation
pip install -r requirements.txt

🤗 Accelerate — это библиотека, которая помогает обучать модели на нескольких GPU/TPU или с использованием смешанной точности. Она автоматически настраивает среду обучения в зависимости от вашего оборудования и окружения. Подробнее можно узнать в Кратком обзоре 🤗 Accelerate.

Инициализируйте среду 🤗 Accelerate (попробуйте включить torch.compile, чтобы значительно ускорить обучение):

accelerate config

Чтобы настроить среду по умолчанию 🤗 Accelerate без выбора каких-либо конфигураций:

accelerate config default

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

from accelerate.utils import write_basic_config

write_basic_config()

Наконец, если вы хотите обучить модель на собственном наборе данных, ознакомьтесь с руководством Create a dataset for training, чтобы узнать, как создать набор данных, совместимый со скриптом обучения.

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

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

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

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

accelerate launch train_lcm_distill_sd_wds.py \
  --mixed_precision="fp16"

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

• --pretrained_teacher_model: путь к предобученной латентной диффузионной модели, используемой в качестве teacher-модели
• --pretrained_vae_model_name_or_path: путь к предобученной VAE; известно, что VAE в SDXL страдает от численной нестабильности, поэтому этот параметр позволяет указать альтернативную VAE (например, VAE) от madebyollin, которая работает в fp16)
• --w_min и --w_max: минимальное и максимальное значения guidance scale для выборки guidance scale
• --num_ddim_timesteps: количество шагов для сэмплирования DDIM
• --loss_type: тип функции потерь (L2 или Huber), используемой для расчёта в процессе дистилляции латентной консистентности; как правило, предпочтение отдаётся функции потерь Huber, поскольку она более устойчива к выбросам
• --huber_c: параметр функции потерь Хабера

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

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

def transform(example):
    image = example["image"]
    image = TF.resize(image, resolution, interpolation=transforms.InterpolationMode.BILINEAR)

    c_top, c_left, _, _ = transforms.RandomCrop.get_params(image, output_size=(resolution, resolution))
    image = TF.crop(image, c_top, c_left, resolution, resolution)
    image = TF.to_tensor(image)
    image = TF.normalize(image, [0.5], [0.5])

    example["image"] = image
    return example

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

processing_pipeline = [
    wds.decode("pil", handler=wds.ignore_and_continue),
    wds.rename(image="jpg;png;jpeg;webp", text="text;txt;caption", handler=wds.warn_and_continue),
    wds.map(filter_keys({"image", "text"})),
    wds.map(transform),
    wds.to_tuple("image", "text"),
]

В функции main() загружаются все необходимые компоненты: планировщик шума, токенизаторы, текстовые энкодеры и VAE. Здесь же загружается teacher UNet, на основе которого можно создать student UNet. Во время обучения student UNet обновляется оптимизатором.

teacher_unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_teacher_model, subfolder="unet", revision=args.teacher_revision
)

unet = UNet2DConditionModel(**teacher_unet.config)
unet.load_state_dict(teacher_unet.state_dict(), strict=False)
unet.train()

Теперь можно создать optimizer для обновления параметров UNet:

optimizer = optimizer_class(
    unet.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Создайте набор данных:

dataset = Text2ImageDataset(
    train_shards_path_or_url=args.train_shards_path_or_url,
    num_train_examples=args.max_train_samples,
    per_gpu_batch_size=args.train_batch_size,
    global_batch_size=args.train_batch_size * accelerator.num_processes,
    num_workers=args.dataloader_num_workers,
    resolution=args.resolution,
    shuffle_buffer_size=1000,
    pin_memory=True,
    persistent_workers=True,
)
train_dataloader = dataset.train_dataloader

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

pred_x_0 = predicted_origin(
    noise_pred,
    start_timesteps,
    noisy_model_input,
    noise_scheduler.config.prediction_type,
    alpha_schedule,
    sigma_schedule,
)

model_pred = c_skip_start * noisy_model_input + c_out_start * pred_x_0

Далее он получает предсказания teacher-модели и предсказания LCM, вычисляет потерю, а затем выполняет обратное распространение ошибки в LCM.

if args.loss_type == "l2":
    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
elif args.loss_type == "huber":
    loss = torch.mean(
        torch.sqrt((model_pred.float() - target.float()) ** 2 + args.huber_c**2) - args.huber_c
    )

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

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

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

В этом руководстве вы используете --train_shards_path_or_url, чтобы указать путь к набору данных Conceptual Captions 12M, хранящемуся на Hub здесь. Установите переменную окружения MODEL_DIR в имя teacher-модели, а OUTPUT_DIR — в каталог, куда вы хотите сохранить модель.

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/saved/model"

accelerate launch train_lcm_distill_sd_wds.py \
    --pretrained_teacher_model=$MODEL_DIR \
    --output_dir=$OUTPUT_DIR \
    --mixed_precision=fp16 \
    --resolution=512 \
    --learning_rate=1e-6 --loss_type="huber" --ema_decay=0.95 --adam_weight_decay=0.0 \
    --max_train_steps=1000 \
    --max_train_samples=4000000 \
    --dataloader_num_workers=8 \
    --train_shards_path_or_url="pipe:curl -L -s https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset/resolve/main/data/{00000..01099}.tar?download=true" \
    --validation_steps=200 \
    --checkpointing_steps=200 --checkpoints_total_limit=10 \
    --train_batch_size=12 \
    --gradient_checkpointing --enable_xformers_memory_efficient_attention \
    --gradient_accumulation_steps=1 \
    --use_8bit_adam \
    --resume_from_checkpoint=latest \
    --report_to=wandb \
    --seed=453645634 \
    --push_to_hub

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

from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
import torch

unet = UNet2DConditionModel.from_pretrained("your-username/your-model", dtype=torch.float16, variant="fp16")
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, dtype=torch.float16, variant="fp16")

pipeline.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "sushi rolls in the form of panda heads, sushi platter"

image = pipeline(prompt, num_inference_steps=4, guidance_scale=1.0).images[0]

LoRA

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

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

Stable Diffusion XL

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

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

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

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

• Узнайте, как использовать LCMs для инференса в задачах генерации изображений по тексту, преобразования изображений и с чекпоинтами LoRA.
• Ознакомьтесь с статьей в блоге «SDXL за 4 шага с помощью Latent Consistency LoRAs», чтобы узнать больше о SDXL LCM-LoRA для сверхбыстрого вывода, сравнении качества, бенчмарках и многом другом.