# Latent Consistency Distillation

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

Если вы обучаете модель на GPU с ограниченным объёмом vRAM, попробуйте включить `gradient_checkpointing`, `gradient_accumulation_steps` и `mixed_precision`, чтобы снизить потребление памяти и ускорить обучение. Вы можете ещё больше сократить использование памяти, включив энергоэффективное внимание с помощью [xFormers](../optimization/attention_backends) и 8-битного оптимизатора [bitsandbytes'](https://github.com/TimDettmers/bitsandbytes).

В этом руководстве мы рассмотрим скрипт [train_lcm_distill_sd_wds.py](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_sd_wds.py), чтобы вы лучше с ним познакомились и могли адаптировать его под свои задачи.

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

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

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

```bash
cd examples/consistency_distillation
pip install -r requirements.txt
```

> [!TIP]
> 🤗 Accelerate — это библиотека, которая помогает обучать модели на нескольких GPU/TPU или с использованием смешанной точности. Она автоматически настраивает среду обучения в зависимости от вашего оборудования и окружения. Подробнее можно узнать в [Кратком обзоре](https://huggingface.co/docs/accelerate/quicktour) 🤗 Accelerate.

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

```bash
accelerate config
```

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

```bash
accelerate config default
```

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

```py
from accelerate.utils import write_basic_config

write_basic_config()
```

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

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

> [!TIP]
> В следующих разделах выделены части скрипта обучения, важные для понимания того, как его модифицировать, но не все аспекты скрипта описаны подробно. Если вы хотите узнать больше, не стесняйтесь прочитать [скрипт](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_sd_wds.py) и сообщить нам, если у вас возникнут вопросы или замечания.

Скрипт обучения предоставляет множество параметров, которые помогают настроить процесс обучения. Все параметры и их описания находятся в функции [`parse_args()`](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L419). Эта функция задаёт значения по умолчанию для каждого параметра, например, размер обучающей выборки и скорость обучения, но при необходимости вы можете указать свои значения в команде запуска обучения.

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

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

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

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

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

Скрипт обучения начинается с создания класса набора данных - [`Text2ImageDataset`](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L141) - для предобработки изображений и формирования обучающего набора данных.

```py
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](https://github.com/webdataset/webdataset) для создания конвейера предобработки, который применяет трансформации и формирует набор данных и загрузчик данных для обучения. Изображения обрабатываются и передаются в цикл обучения без необходимости предварительно скачивать весь набор данных.

```py
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()`](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L768) загружаются все необходимые компоненты: планировщик шума, токенизаторы, текстовые энкодеры и VAE. Здесь же загружается teacher UNet, на основе которого можно создать student UNet. Во время обучения student UNet обновляется оптимизатором.

```py
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](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L979) для обновления параметров UNet:

```py
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,
)
```

Создайте [набор данных](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L994):

```py
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
```

Далее можно настроить [цикл обучения](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L1049) и реализовать метод дистилляции латентной согласованности (подробности см. в Алгоритме 1 статьи). В этой части скрипта добавляется шум к латентам, выполняется выборка и создание эмбеддинга масштаба гайданса, а также предсказывается исходное изображение по шуму.

```py
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-модели](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L1172) и [предсказания LCM](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L1209), вычисляет потерю, а затем выполняет обратное распространение ошибки в LCM.

```py
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
    )
```

Если вы хотите узнать больше о том, как работает цикл обучения, ознакомьтесь с [руководством «Понимание конвейеров, моделей и планировщиков»](../using-diffusers/write_own_pipeline), в котором подробно разбирается базовый паттерн процесса устранения шума.

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

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

В этом руководстве вы используете `--train_shards_path_or_url`, чтобы указать путь к набору данных [Conceptual Captions 12M](https://github.com/google-research-datasets/conceptual-12m), хранящемуся на Hub [здесь](https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset). Установите переменную окружения `MODEL_DIR` в имя teacher-модели, а `OUTPUT_DIR` — в каталог, куда вы хотите сохранить модель.

```bash
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 для инференса.

```py
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](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py) или [train_lcm_distill_lora_sdxl.wds.py](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py).

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

## Stable Diffusion XL

Stable Diffusion XL (SDXL) — это мощная модель text-to-image, генерирующая изображения высокого разрешения, которая добавляет второй text-encoder в свою архитектуру. Используйте скрипт [train_lcm_distill_sdxl_wds.py](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py) для обучения модели SDXL с LoRA.

Скрипт обучения SDXL подробно описан в руководстве [обучение SDXL](sdxl).

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

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

- Узнайте, как использовать [LCMs для инференса](../api/pipelines/latent_consistency_models) в задачах генерации изображений по тексту, преобразования изображений и с чекпоинтами LoRA.
- Ознакомьтесь с [статьей в блоге «SDXL за 4 шага с помощью Latent Consistency LoRAs»](https://huggingface.co/blog/lcm_lora), чтобы узнать больше о SDXL LCM-LoRA для сверхбыстрого вывода, сравнении качества, бенчмарках и многом другом.