LlamaFactory
算力平台:

DreamBooth ​

DreamBooth 是一种训练技术，通过仅使用少量主题或风格的图像来更新整个扩散模型。它通过将提示中的一个特殊词与示例图像关联起来来实现这一点。

如果你在 vRAM 有限的 GPU 上进行训练，建议在训练命令中启用 gradient_checkpointing 和 mixed_precision 参数。你还可以通过使用 xFormers 的内存高效注意力机制来减少内存占用。JAX/Flax 训练也支持在 TPUs 和 GPUs 上高效训练，但不支持梯度检查点或 xFormers。如果你希望使用 Flax 进行更快的训练，建议使用内存大于 30GB 的 GPU。

本指南将探讨 train_dreambooth.py 脚本，帮助你更熟悉它，并了解如何根据自己的需求进行调整。

在运行脚本之前，请确保从源代码安装库：

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

导航到包含训练脚本的示例文件夹，并安装你所使用脚本所需的依赖项：

初始化一个 🤗 Accelerate 环境：

bash
accelerate config

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

bash
accelerate config default

或者，如果你的环境不支持交互式 shell，比如笔记本，你可以使用：

py
from accelerate.utils import write_basic_config

write_basic_config()

最后，如果你想在自己的数据集上训练模型，请参阅创建用于训练的数据集指南，了解如何创建与训练脚本兼容的数据集。

脚本参数 ​

训练脚本提供了许多参数，用于自定义你的训练过程。所有参数及其描述都在 parse_args() 函数中。这些参数已经设置了默认值，通常情况下可以直接使用，但你也可以在训练命令中设置自己的值。

例如，要以 bf16 格式进行训练：

bash
accelerate launch train_dreambooth.py \
    --mixed_precision="bf16"

一些基本且重要的参数需要了解和指定：

  • --pretrained_model_name_or_path：Hub 上的模型名称或预训练模型的本地路径
  • --instance_data_dir：包含训练数据集（示例图像）的文件夹路径
  • --instance_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：

bash
accelerate launch train_dreambooth.py \
  --snr_gamma=5.0

先验保留损失 ​

先验保留损失是一种方法，它使用模型自身生成的样本帮助模型学习如何生成更多样化的图像。因为这些生成的样本图像与你提供的图像属于同一类别，所以它们有助于模型保留对类别的学习，并利用已知的类别知识来创建新的组合。

  • --with_prior_preservation: 是否使用先验保留损失
  • --prior_loss_weight: 控制先验保留损失对模型的影响
  • --class_data_dir: 包含生成的类别样本图像的文件夹路径
  • --class_prompt: 描述生成样本图像类别的文本提示
bash
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。如果你有必要的硬件，那么训练文本编码器会产生更好的结果，尤其是在生成人脸图像时。启用此选项的方法如下：

bash
accelerate launch train_dreambooth.py \
  --train_text_encoder

训练脚本 ​

DreamBooth 配备了自己的数据集类：

如果你启用了 先验保留损失，类别图像将在这里生成：

py
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() 函数，该函数负责设置训练数据集和训练循环本身。脚本加载了 分词器、调度器和模型：

py
# 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 了。

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

最后，训练循环 负责剩余的步骤，例如将图像转换为潜在空间、向输入添加噪声、预测噪声残差以及计算损失。

如果你想了解更多关于训练循环的工作原理，可以查看 理解管道、模型和调度器 教程，该教程详细介绍了去噪过程的基本模式。

启动脚本 ​

你现在可以启动训练脚本了！🚀

在本指南中，你将下载一些 狗 的图像并将其存储在一个目录中。但请记住，你也可以创建并使用自己的数据集（参见 为训练创建数据集 指南）。

py
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 作为特殊词来关联训练。

如果你有兴趣跟踪训练过程，可以在训练进行时定期保存生成的图像。在训练命令中添加以下参数：

bash
--validation_prompt="a photo of a sks dog"
--num_validation_images=4
--validation_steps=100

在你启动脚本之前还有一件事！根据你拥有的 GPU，你可能需要启用某些优化来训练 DreamBooth。

一旦训练完成，你就可以使用新训练的模型进行推理！

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 权重时，还必须更新管道的调度器配置。

py
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 阶段模型需要额外的验证图像来进行放大。你可以下载并使用训练图像的缩小版本来完成这一任务。

py
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，提前计算文本嵌入以节省内存，因为 [~transformers.T5Model] 可能会占用大量内存
  • --tokenizer_max_length=77，你可以使用更长的默认文本长度，但默认的模型编码过程使用较短的文本长度
  • --text_encoder_use_attention_mask，将注意力掩码传递给文本编码器

训练技巧 ​

训练 DeepFloyd IF 模型可能会有挑战，但以下是一些我们发现有用的技巧：

  • LoRA 对于训练第一阶段模型已经足够，因为模型的低分辨率使得表示更精细的细节变得困难。
  • 对于常见或简单的对象，你不一定需要微调上采样器。确保传递给上采样器的提示词已调整，以移除实例提示词中的新标记。例如，如果你的第一阶段提示词是 "a sks dog"，那么你的第二阶段提示词应该是 "a dog"。
  • 对于更精细的细节，如面部，完全训练第二阶段上采样器比使用 LoRA 训练第二阶段模型更好。使用较低的学习率和较大的批量也有帮助。
  • 训练第二阶段模型时应使用较低的学习率。
  • 使用 [DDPMScheduler] 比训练脚本中使用的 DPMSolver 更好。

下一步 ​

恭喜你训练了你的 DreamBooth 模型！要了解如何使用你的新模型，以下指南可能会有所帮助：

  • 如果你使用 LoRA 训练了模型，可以学习如何 加载 DreamBooth 模型进行推理。

Copyright © 2024-2025 llamafactory.cn