Diffusers documentation

文本反转

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

文本反转

Textual Inversion 是一种训练技术,用于通过少量示例图像来个性化图像生成模型,使其学习您希望它掌握的内容。该技术通过学习并更新文本嵌入(新嵌入与提示词中必须使用的特殊单词相关联)来匹配您提供的示例图像。

有关使用训练好的嵌入进行推理,请参阅 文本反转推理。

如果用于训练的 GPU 显存有限,建议在训练命令中启用 gradient_checkpointing 和 mixed_precision 参数。此外,你还可以使用 xFormers 提供的内存高效注意力机制来降低内存占用。

本指南将深入讲解 textual_inversion.py 脚本,帮助你更熟悉该脚本,并了解如何将其适配到自身的使用场景中。

在运行脚本之前,请确保从源代码安装该库:

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

导航至包含训练脚本的示例文件夹,并安装您所使用的脚本所需的依赖项:

cd examples/textual_inversion
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()

最后,如果您想在自己的数据集上训练模型,请查看创建训练数据集指南,了解如何创建可与训练脚本配合使用的数据集。

以下各部分重点介绍了训练脚本中对于理解如何修改脚本至关重要的内容,但并未详细涵盖脚本的每个方面。如果您希望了解更多,请随时阅读脚本,并告诉我们您是否有任何疑问或顾虑。

脚本参数

训练脚本包含许多参数,可帮助您根据需求定制训练过程。所有参数及其说明均列于 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 用于创建数据集。您可以自定义图像尺寸、占位符 token、插值方法、是否裁剪图像等。如果需要更改数据集的创建方式,可以修改 TextualInversionDataset。

接下来,你可以在 main() 函数中找到数据集预处理代码和训练循环。

脚本首先加载 tokenizer、scheduler 和模型:

# 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:特殊占位符 token
  • type_of_concept.txt:您正在训练的概念类型(“object”或“style”)

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

在启动脚本之前,还有一件事。如果您希望跟踪训练过程,可以在训练进行期间定期保存生成的图像。请将以下参数添加到训练命令中:

--validation_prompt="A <cat-toy> train"
--num_validation_images=4
--validation_steps=100
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

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

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
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 模型!🎉 如需了解如何使用新模型,以下指南可能对您有所帮助:

Update on GitHub