Diffusers 文档
DreamBooth
并获得增强的文档体验
开始使用
DreamBooth
DreamBooth 是一种训练技术，通过仅在少量主题或风格的图像上进行训练来更新整个扩散模型。它的工作原理是将 prompt 中的一个特殊词语与示例图像关联起来。
如果您在 vRAM 有限的 GPU 上进行训练，您应该尝试在训练命令中启用 gradient_checkpointing 和 mixed_precision 参数。您还可以通过使用具有 xFormers 的内存高效注意力来减少内存占用。JAX/Flax 训练也支持在 TPU 和 GPU 上进行高效训练，但它不支持梯度检查点或 xFormers。如果您想使用 Flax 更快地进行训练，您应该拥有一个内存 >30GB 的 GPU。
本指南将探讨 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
或者，如果您的环境不支持交互式 shell，例如 notebook，您可以使用
from accelerate.utils import write_basic_config
write_basic_config()最后，如果您想在自己的数据集上训练模型，请查看 创建用于训练的数据集 指南，了解如何创建与训练脚本配合使用的数据集。
以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分，但它没有详细介绍脚本的每个方面。如果您有兴趣了解更多信息，请随时通读 脚本，如果您有任何问题或疑虑，请告知我们。
脚本参数
DreamBooth 对训练超参数非常敏感，并且很容易过拟合。请阅读 使用 🧨 Diffusers 通过 Dreambooth 训练 Stable Diffusion 博客文章，了解针对不同主题的推荐设置，以帮助您选择合适的超参数。
训练脚本提供了许多参数来自定义您的训练运行。所有参数及其描述都可以在 parse_args() 函数中找到。这些参数设置了默认值，这些默认值应该可以很好地开箱即用，但如果您愿意，也可以在训练命令中设置您自己的值。
例如，以 bf16 格式进行训练
accelerate launch train_dreambooth.py \
--mixed_precision="bf16"一些基本且重要的参数需要了解和指定：
--pretrained_model_name_or_path：Hub 上的模型名称或预训练模型的本地路径--instance_data_dir：包含训练数据集（示例图像）的文件夹路径--instance_prompt：包含示例图像特殊词语的文本 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 支持，在 Flax 训练脚本中不可用。
添加 --snr_gamma 参数并将其设置为推荐值 5.0
accelerate launch train_dreambooth.py \ --snr_gamma=5.0
先验保留损失
先验保留损失是一种方法，它使用模型自身生成的样本来帮助其学习如何生成更多样化的图像。由于这些生成的样本图像与您提供的图像属于同一类别，因此它们有助于模型保留其已学到的关于该类别的知识，以及如何使用其已知的关于该类别的知识来创建新的构图。
--with_prior_preservation：是否使用先验保留损失--prior_loss_weight：控制先验保留损失对模型的影响--class_data_dir：包含生成的类样本图像的文件夹路径--class_prompt：描述生成的样本图像类别的文本 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。这需要额外的内存，您需要至少 24GB vRAM 的 GPU。如果您有必要的硬件，那么训练文本编码器会产生更好的结果，尤其是在生成人脸图像时。通过以下方式启用此选项：
accelerate launch train_dreambooth.py \ --train_text_encoder
训练脚本
DreamBooth 自带数据集类
DreamBoothDataset：预处理图像和类别图像，并对 prompt 进行分词以进行训练PromptDataset：生成 prompt 嵌入以生成类别图像
如果您启用了 先验保留损失，则在此处生成类别图像
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
)然后，是时候从 DreamBoothDataset 创建训练数据集 和 DataLoader 了
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,
)最后，训练循环 处理剩余步骤，例如将图像转换为潜在空间、向输入添加噪声、预测噪声残差以及计算损失。
如果您想了解有关训练循环如何工作的更多信息，请查看 理解 pipelines、模型和调度器 教程，其中分解了去噪过程的基本模式。
启动脚本
您现在可以启动训练脚本了！ 🚀
在本指南中，您将下载一些狗的图像并将它们存储在一个目录中。但请记住，如果您愿意，您可以创建和使用您自己的数据集（请参阅“创建用于训练的数据集”指南）。
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 上的模型 ID 或本地模型的路径，INSTANCE_DIR 设置为您刚刚下载狗图像的路径，OUTPUT_DIR 设置为您想要保存模型的路径。您将使用 sks 作为特殊词语来将训练与之关联。
如果您有兴趣跟进训练过程，您可以定期保存随着训练进展而生成的图像。将以下参数添加到训练命令中
--validation_prompt="a photo of a sks dog"
--num_validation_images=4
--validation_steps=100在启动脚本之前还有一件事！根据您拥有的 GPU，您可能需要启用某些优化来训练 DreamBooth。
在 16GB GPU 上，您可以使用 bitsandbytes 8 位优化器和梯度检查点来帮助您训练 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", torch_dtype=torch.float16, use_safetensors=True).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")LoRA
LoRA 是一种训练技术，用于显著减少可训练参数的数量。因此，训练速度更快，并且更容易存储结果权重，因为它们小得多（约 100MB）。使用 train_dreambooth_lora.py 脚本来使用 LoRA 进行训练。
LoRA 训练脚本在 LoRA 训练指南中进行了更详细的讨论。
Stable Diffusion XL
Stable Diffusion XL (SDXL) 是一款强大的文本到图像模型，可以生成高分辨率图像，并在其架构中添加了第二个文本编码器。使用 train_dreambooth_lora_sdxl.py 脚本来使用 LoRA 训练 SDXL 模型。
SDXL 训练脚本在 SDXL 训练指南中进行了更详细的讨论。
DeepFloyd IF
DeepFloyd IF 是一个具有三个阶段的级联像素扩散模型。第一阶段生成基础图像，第二和第三阶段逐步将基础图像放大为高分辨率 1024x1024 图像。使用 train_dreambooth_lora.py 或 train_dreambooth.py 脚本来使用 LoRA 或完整模型训练 DeepFloyd IF 模型。
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",
)下面的代码示例简要概述了如何结合 DreamBooth 和 LoRA 来训练 DeepFloyd IF 模型。需要注意的一些重要参数是
--resolution=64，需要更小的分辨率，因为 DeepFloyd IF 是一个像素扩散模型，并且为了处理未压缩的像素，输入图像必须更小--pre_compute_text_embeddings，提前计算文本嵌入以节省内存，因为 T5Model 可能会占用大量内存--tokenizer_max_length=77，您可以将更长的默认文本长度与 T5 用作文本编码器，但默认模型编码过程使用较短的文本长度--text_encoder_use_attention_mask，将注意力掩码传递给文本编码器
使用 LoRA 和 DreamBooth 训练 DeepFloyd IF 的阶段 1 需要约 28GB 内存。
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 对于训练阶段 1 模型已经足够，因为无论如何，模型的分辨率较低使得表示更精细的细节变得困难。
- 对于常见的或简单的对象，您不一定需要微调放大器。确保传递给放大器的提示已调整为从实例提示中删除新 token。例如，如果您的阶段 1 提示是 “a sks dog”，那么您的阶段 2 提示应该是 “a dog”。
- 对于更精细的细节，如面部，完全训练阶段 2 放大器比使用 LoRA 训练阶段 2 模型更好。使用较低的学习率和较大的批次大小也很有帮助。
- 应使用较低的学习率来训练阶段 2 模型。
- DDPMScheduler 比训练脚本中使用的 DPMSolver 效果更好。
下一步
祝贺您训练了您的 DreamBooth 模型！要了解有关如何使用您的新模型的更多信息，以下指南可能会有所帮助
- 如果您使用 LoRA 训练了模型，请了解如何加载 DreamBooth 模型以进行推理。