Hugging Face 的 logo

Diffusers 文档

文本反转

Hugging Face's logo
加入 Hugging Face 社区

并获取增强文档体验

开始使用

文本反转

文本反转 是一种训练技术，只需使用您想要其学习的内容的几个示例图像即可个性化图像生成模型。此技术通过学习和更新文本嵌入（新的嵌入与您必须在提示中使用的特殊单词绑定）来匹配您提供的示例图像。

如果您在 vRAM 有限的 GPU 上进行训练，则应尝试在训练命令中启用 gradient_checkpointing 和 mixed_precision 参数。您还可以通过使用具有 xFormers 的内存高效注意力来减少内存占用。JAX/Flax 训练也支持在 TPU 和 GPU 上进行高效训练，但它不支持梯度检查点或 xFormers。在与 PyTorch 相同的配置和设置下，Flax 训练脚本的速度至少可以提高约 70%！

本指南将探讨 textual_inversion.py 脚本，以帮助您对其更加熟悉，以及如何将其调整以适应您自己的用例。

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

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

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

PyTorch
Flax
cd examples/textual_inversion
pip install -r requirements.txt

🤗 Accelerate 是一个帮助您在多个 GPU/TPU 或混合精度上进行训练的库。它将根据您的硬件和环境自动配置您的训练设置。查看 🤗 Accelerate 的 快速入门 了解更多信息。

初始化 🤗 Accelerate 环境

accelerate config

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

accelerate config default

或者，如果您的环境不支持交互式 shell（例如笔记本），则可以使用

from accelerate.utils import write_basic_config

write_basic_config()

最后，如果您想在自己的数据集上训练模型，请查看 创建用于训练的数据集 指南，了解如何创建适用于训练脚本的数据集。

以下部分重点介绍了训练脚本中一些了解如何修改它的重要部分，但它没有详细介绍脚本的每个方面。如果您有兴趣了解更多信息，请随时通读 脚本 并告诉我们您有任何疑问或疑虑。

脚本参数

训练脚本有许多参数可以帮助您根据需要调整训练运行。所有参数及其描述都列在 parse_args() 函数中。在适用的情况下，Diffusers 为每个参数提供默认值，例如训练批大小和学习率，但如果您愿意，可以在训练命令中更改这些值。

例如，要将梯度累积步骤的数量增加到默认值 1 以上

accelerate launch textual_inversion.py \
  --gradient_accumulation_steps=4

其他一些基本且重要的参数包括

  • --pretrained_model_name_or_path：Hub 上模型的名称或预训练模型的本地路径
  • --train_data_dir：包含训练数据集（示例图像）的文件夹的路径
  • --output_dir：保存训练模型的位置
  • --push_to_hub：是否将训练模型推送到 Hub
  • --checkpointing_steps：在模型训练时保存检查点的频率；如果由于某种原因训练中断，您可以通过将 --resume_from_checkpoint 添加到训练命令来从该检查点继续训练
  • --num_vectors：学习嵌入向量的数量；增加此参数有助于模型学习得更好，但会增加训练成本
  • --placeholder_token：将学习到的嵌入绑定到的特殊单词（您必须在推理提示中使用该单词）
  • --initializer_token：大致描述您尝试训练的对象或样式的单个单词
  • --learnable_property：您是否正在训练模型以学习新的“样式”（例如，梵高的绘画风格）或“对象”（例如，您的狗）

训练脚本

与其他一些训练脚本不同，textual_inversion.py 有一个自定义数据集类 TextualInversionDataset 用于创建数据集。您可以自定义图像大小、占位符标记、插值方法、是否裁剪图像等等。如果您需要更改数据集的创建方式，可以修改 TextualInversionDataset。

接下来，您将在 main() 函数中找到数据集预处理代码和训练循环。

该脚本首先加载 标记器、调度器和模型

# Load tokenizer
if args.tokenizer_name:
    tokenizer = CLIPTokenizer.from_pretrained(args.tokenizer_name)
elif args.pretrained_model_name_or_path:
    tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer")

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = CLIPTextModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

接下来，将特殊的 占位符标记 添加到标记器中，并重新调整嵌入以考虑新标记。

然后，脚本 从 TextualInversionDataset 创建数据集

train_dataset = TextualInversionDataset(
    data_root=args.train_data_dir,
    tokenizer=tokenizer,
    size=args.resolution,
    placeholder_token=(" ".join(tokenizer.convert_ids_to_tokens(placeholder_token_ids))),
    repeats=args.repeats,
    learnable_property=args.learnable_property,
    center_crop=args.center_crop,
    set="train",
)
train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.dataloader_num_workers
)

最后，训练循环 处理从预测噪声残差到更新特殊占位符标记的嵌入权重的所有其他操作。

如果您想了解有关训练循环如何工作的更多信息，请查看 理解管道、模型和调度器 教程，该教程分解了去噪过程的基本模式。

启动脚本

完成所有更改或对默认配置感到满意后，您就可以启动训练脚本了！🚀

在本指南中，您将下载一些猫咪玩具的图像并将其存储在一个目录中。但请记住，如果您愿意，可以创建和使用您自己的数据集（请参阅创建训练数据集指南）。

from huggingface_hub import snapshot_download

local_dir = "./cat"
snapshot_download(
    "diffusers/cat_toy_example", local_dir=local_dir, repo_type="dataset", ignore_patterns=".gitattributes"
)

将环境变量MODEL_NAME设置为 Hub 上的模型 ID 或本地模型的路径，并将DATA_DIR设置为您刚刚下载猫咪图像的路径。脚本会创建以下文件并将其保存到您的存储库中

  • learned_embeds.bin：与您的示例图像对应的学习到的嵌入向量
  • token_identifier.txt：特殊的占位符标记
  • type_of_concept.txt：您正在训练的概念类型（“object”或“style”）

完整的训练过程在一个 V100 GPU 上大约需要 1 小时。

在启动脚本之前还有一件事。如果您有兴趣关注训练过程，可以在训练过程中定期保存生成的图像。将以下参数添加到训练命令中

--validation_prompt="A <cat-toy> train"
--num_validation_images=4
--validation_steps=100
PyTorch
Flax
export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export DATA_DIR="./cat"

accelerate launch textual_inversion.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATA_DIR \
  --learnable_property="object" \
  --placeholder_token="<cat-toy>" \
  --initializer_token="toy" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=3000 \
  --learning_rate=5.0e-04 \
  --scale_lr \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --output_dir="textual_inversion_cat" \
  --push_to_hub

训练完成后，您可以像这样使用新训练的模型进行推理

PyTorch
Flax
from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
pipeline.load_textual_inversion("sd-concepts-library/cat-toy")
image = pipeline("A <cat-toy> train", num_inference_steps=50).images[0]
image.save("cat-train.png")

下一步

恭喜您训练了自己的文本反转模型！🎉 若要详细了解如何使用您的新模型，以下指南可能会有所帮助

< > GitHub 更新

←InstructPix2Pix DreamBooth→