# 潜在一致性蒸馏

[潜在一致性模型（LCMs）](https://hf.co/papers/2310.04378)能够在极少步数内生成高质量图像，这是一项重大突破，因为许多流水线至少需要 25 步以上。LCMs 通过对任意 Stable Diffusion 模型应用潜在一致性蒸馏方法而生成。该方法通过在潜在空间中应用*单阶段引导蒸馏*，并结合*跳步*方法以一致地跳过时间步来加速蒸馏过程（更多细节请参阅论文的第 4.1、4.2 和 4.3 节）。

如果您在 vRAM 有限的 GPU 上进行训练，请尝试启用 `gradient_checkpointing`、`gradient_accumulation_steps` 和 `mixed_precision`，以降低内存占用并加快训练速度。您还可以通过启用 [xFormers](../optimization/attention_backends) 和 [bitsandbytes'](https://github.com/TimDettmers/bitsandbytes) 8 位优化器来实现内存高效注意力机制，从而进一步降低内存占用。

本指南将深入讲解 [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 上训练或使用混合精度训练的库。它会根据你的硬件和环境自动配置训练设置。请查看 🤗 Accelerate 的 [快速入门](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。

初始化一个 🤗 Accelerate 环境（建议启用 `torch.compile` 以显著加快训练速度）：

```bash
accelerate config
```

要设置一个默认的 🤗 Accelerate 环境，而不选择任何配置：

```bash
accelerate config default
```

或者，如果你的环境不支持交互式 shell（例如 notebook），你可以使用：

```py
from accelerate.utils import write_basic_config

write_basic_config()
```

最后，如果您想在自己的数据集上训练模型，请查看[创建训练数据集](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"
```

大多数参数与[文本到图像](text2image#script-parameters)训练指南中的参数相同，因此本指南将重点介绍与潜在一致性蒸馏相关的参数。

- `--pretrained_teacher_model`：用作教师模型的预训练潜在扩散模型的路径
- `--pretrained_vae_model_name_or_path`：预训练 VAE 的路径；SDXL 的 VAE 已知存在数值不稳定性问题，因此该参数允许你指定一个替代的 VAE（例如 madebyollin 提供的、支持 fp16 的 [VAE](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix)）
- `--w_min` 和 `--w_max`：用于 guidance scale 采样的最小和最大 guidance scale 值
- `--num_ddim_timesteps`：DDIM 采样的时间步数
- `--loss_type`：用于计算潜在一致性蒸馏的损失类型（L2 或 Huber）；通常更推荐使用 Huber 损失，因为它对异常值具有更强的鲁棒性
- `--huber_c`：Huber 损失参数

## 训练脚本

训练脚本首先创建一个数据集类 - [`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。教师 UNet 也会在此处加载，随后可以基于教师 UNet 创建学生 UNet。在训练过程中，优化器会更新学生 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,
)
```

创建 [dataset](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
```

接下来，它会获取[教师模型的预测](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` 来指定存储在 Hub 上的 [Conceptual Captions 12M](https://github.com/google-research-datasets/conceptual-12m) 数据集的路径，该数据集位于 [此处](https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset)。将 `MODEL_DIR` 环境变量设置为教师模型的名称，并将 `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 是一种训练技术，可大幅减少可训练参数的数量。因此，训练速度更快，且由于生成的权重文件小得多（约 100MB），存储也更为便捷。请使用 [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 训练](lora) 指南中有更详细的介绍。

## Stable Diffusion XL

Stable Diffusion XL (SDXL) 是一个强大的文本生成图像模型，能够生成高分辨率图像，并在其架构中增加了第二个文本编码器。使用 [train_lcm_distill_sdxl_wds.py](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py) 脚本通过 LoRA 训练 SDXL 模型。

SDXL 训练脚本在 [SDXL 训练](sdxl) 指南中有更详细的介绍。

## 后续步骤

恭喜您成功蒸馏了一个 LCM 模型！如需进一步了解 LCM，以下内容可能对您有所帮助：

- 了解如何使用 [LCM 进行推理](../api/pipelines/latent_consistency_models)，包括文本生成图像、图像生成图像，以及结合 LoRA 检查点。
- 阅读[使用潜在一致性 LoRA 在 4 步内运行 SDXL](https://huggingface.co/blog/lcm_lora)博客文章，以了解有关 SDXL LCM-LoRA 的更多信息，包括超快推理、质量对比、基准测试等。