潜在一致性蒸馏
潜在一致性模型 (LCM) 能够在短短几步内生成高质量图像，这是一个巨大的进步，因为许多管道需要至少 25 步以上才能生成图像。LCM 是通过将潜在一致性蒸馏方法应用于任何 Stable Diffusion 模型而产生的。此方法通过在潜在空间中应用 *单阶段引导蒸馏*，并将 *跳步* 方法整合进来以一致地跳过时间步以加速蒸馏过程来实现（有关更多详细信息，请参阅论文的第 4.1、4.2 和 4.3 节）。
如果你在 vRAM 有限的 GPU 上进行训练，尝试启用 gradient_checkpointing、gradient_accumulation_steps 和 mixed_precision 来减少内存使用量并加速训练。你可以通过启用具有 xFormers 和 bitsandbytes 的 8 位优化器来减少内存使用量。
本指南将探讨 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
或者，如果你的环境不支持交互式 shell（如笔记本），你可以使用
from accelerate.utils import write_basic_config
write_basic_config()最后，如果你想在自己的数据集上训练模型，请查看 创建训练数据集 指南，了解如何创建适用于训练脚本的数据集。
脚本参数
以下部分重点介绍了训练脚本中的一些重要部分，有助于你理解如何修改它，但它并不详细涵盖脚本的每个方面。如果你有兴趣了解更多信息，请随时阅读 脚本 ，如果有任何问题或疑虑，请告知我们。
训练脚本提供了许多参数，以帮助你自定义训练运行。所有参数及其描述都可以在 parse_args() 函数中找到。此函数为每个参数提供了默认值（例如训练批次大小和学习率），但你也可以在训练命令中设置自己的值（如果你愿意）。
例如，要使用 fp16 格式的混合精度来加速训练，请将 --mixed_precision 参数添加到训练命令中
accelerate launch train_lcm_distill_sd_wds.py \
--mixed_precision="fp16"大多数参数与 文本到图像 训练指南中的参数相同，因此你将重点关注本指南中与潜在一致性蒸馏相关的参数。
--pretrained_teacher_model: 用于作为教师模型的预训练潜在扩散模型的路径--pretrained_vae_model_name_or_path: 预训练 VAE 的路径；已知 SDXL VAE 会出现数值不稳定现象，因此此参数允许你指定一个替代 VAE（例如由 madebyollin 提供的 VAE ，它可以在 fp16 中运行）--w_min和--w_max: 指导尺度采样的最小和最大引导尺度值--num_ddim_timesteps: DDIM 采样的时间步数--loss_type: 用于计算潜在一致性蒸馏的损失类型（L2 或 Huber）；Huber 损失通常更受欢迎，因为它对异常值更稳健--huber_c: Huber 损失参数
训练脚本
训练脚本首先创建一个数据集类 - 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。教师 UNet 也在这里加载，然后你可以从教师 UNet 创建一个学生 UNet。学生 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()现在你可以创建 优化器 来更新 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它获取了 教师模型预测 和 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 来指定存储在 Hub 上的 概念字幕 12M 数据集的路径 这里。将 MODEL_DIR 环境变量设置为教师模型的名称，并将 OUTPUT_DIR 设置为要保存模型的位置。
export MODEL_DIR="runwayml/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", torch_dtype=torch.float16, variant="fp16")
pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", unet=unet, torch_dtype=torch.float16, variant="fp16")
pipeline.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipeline.to("cuda")
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 MB）。使用 train_lcm_distill_lora_sd_wds.py 或 train_lcm_distill_lora_sdxl.wds.py 脚本使用 LoRA 进行训练。
LoRA 训练脚本在 LoRA 训练 指南中进行了更详细的讨论。
稳定扩散 XL
稳定扩散 XL (SDXL) 是一种功能强大的文本到图像模型，它生成高分辨率图像，并在其架构中添加了第二个文本编码器。使用 train_lcm_distill_sdxl_wds.py 脚本使用 LoRA 训练 SDXL 模型。
SDXL 训练脚本在 SDXL 训练 指南中进行了更详细的讨论。
下一步
恭喜您蒸馏了 LCM 模型！要了解有关 LCM 的更多信息，以下内容可能会有所帮助
- 了解如何使用 用于推断的 LCM 进行文本到图像、图像到图像以及使用 LoRA 检查点。
- 阅读 使用潜在一致性 LoRA 进行 4 步 SDXL 博客文章，以了解有关 SDXL LCM-LoRA 的更多信息，以实现超快推断、质量比较、基准测试等等。