Diffusers 文档

文本反转 (Textual Inversion)

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

Textual Inversion

Textual Inversion 是一种训练技术,只需几张你想要学习的示例图片,即可个性化图像生成模型。该技术通过学习和更新文本嵌入(新的嵌入与你必须在提示词中使用的特殊词语相关联)来匹配你提供的示例图片。

如果你在 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:你正在训练的概念类型(“对象”或“风格”)

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

在启动脚本之前,还有一件事。如果你有兴趣跟踪训练过程,可以在训练进行时定期保存生成的图片。在训练命令中添加以下参数:

--validation_prompt="A <cat-toy> train"
--num_validation_images=4
--validation_steps=100
PyTorch
Flax
export MODEL_NAME="stable-diffusion-v1-5/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("stable-diffusion-v1-5/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")

后续步骤

恭喜你训练了自己的 Textual Inversion 模型!🎉 要了解更多关于如何使用新模型的信息,以下指南可能会有所帮助:

< > 在 GitHub 上更新