自定义扩散
自定义扩散 是一种用于个性化图像生成模型的训练技术。与文本反转(Textual Inversion)、DreamBooth 和 LoRA 一样,自定义扩散只需要少数(约 4-5 张)示例图像。该技术通过仅训练交叉注意力层中的权重来工作,并使用一个特殊词汇来表示新学习的概念。自定义扩散的独特之处在于它还可以同时学习多个概念。
如果你在 vRAM 有限的 GPU 上进行训练,建议启用 xFormers,使用 --enable_xformers_memory_efficient_attention
以更快的训练速度和更低的 vRAM 需求(16GB)。为了节省更多内存,可以在训练参数中添加 --set_grads_to_none
,将梯度设置为 None
而不是零(此选项可能会导致一些问题,如果遇到问题,可以尝试移除该参数)。
本指南将探讨 train_custom_diffusion.py 脚本,帮助你更熟悉它,并了解如何根据自己的需求进行调整。
在运行脚本之前,请确保从源代码安装库:
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
导航到包含训练脚本的示例文件夹并安装所需的依赖项:
cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval
初始化一个 🤗 Accelerate 环境:
accelerate config
要设置一个默认的 🤗 Accelerate 环境而不选择任何配置:
accelerate config default
或者,如果你的环境不支持交互式 shell,比如笔记本,你可以使用:
from accelerate.utils import write_basic_config
write_basic_config()
最后,如果你想在自己的数据集上训练模型,请参阅创建用于训练的数据集指南,了解如何创建与训练脚本兼容的数据集。
脚本参数
训练脚本包含所有帮助你自定义训练运行的参数。这些参数位于 parse_args()
函数中。该函数带有默认值,但你也可以在训练命令中设置自己的值。
例如,要更改输入图像的分辨率:
accelerate launch train_custom_diffusion.py \
--resolution=256
许多基本参数在 DreamBooth 训练指南中已有描述,因此本指南将重点介绍 Custom Diffusion 独有的参数:
--freeze_model
: 冻结交叉注意力层中的键和值参数;默认值为crossattn_kv
,但你可以将其设置为crossattn
以训练交叉注意力层中的所有参数--concepts_list
: 若要学习多个概念,请提供一个包含这些概念的 JSON 文件的路径--modifier_token
: 用于表示学习到的概念的特殊单词--initializer_token
: 用于初始化modifier_token
嵌入的特殊单词
先验保留损失
先验保留损失是一种方法,它使用模型自动生成的样本帮助模型学习如何生成更多样化的图像。因为这些生成的样本图像与你提供的图像属于同一类别,所以它们有助于模型保留已学的类别知识,并利用已知的类别知识创建新的组合。
先验保留损失的许多参数在 DreamBooth 训练指南中已有描述。
正则化
Custom Diffusion 包括使用一小部分真实图像训练目标图像以防止过拟合。可以想象,当你只在几张图像上进行训练时,这很容易做到!使用 clip_retrieval
下载 200 张真实图像。class_prompt
应与目标图像属于同一类别。这些图像存储在 class_data_dir
中。
python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200
要启用正则化,请添加以下参数:
--with_prior_preservation
: 是否使用先验保留损失--prior_loss_weight
: 控制先验保留损失对模型的影响--real_prior
: 是否使用一小部分真实图像来防止过拟合
accelerate launch train_custom_diffusion.py \
--with_prior_preservation \
--prior_loss_weight=1.0 \
--class_data_dir="./real_reg/samples_cat" \
--class_prompt="cat" \
--real_prior=True \
训练脚本
自定义扩散训练脚本包含两个数据集类:
CustomDiffusionDataset
: 预处理用于训练的图像、类别图像和提示PromptDataset
: 准备生成类别图像的提示
接下来,将 modifier_token
添加到分词器,转换为 token id,并调整 token 嵌入的大小以适应新的 modifier_token
。然后,使用 initializer_token
的嵌入初始化 modifier_token
的嵌入。文本编码器中的所有参数都被冻结,除了 token 嵌入,因为这是模型试图学习与概念关联的部分。
params_to_freeze = itertools.chain(
text_encoder.text_model.encoder.parameters(),
text_encoder.text_model.final_layer_norm.parameters(),
text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)
现在你需要将自定义扩散权重添加到注意力层。这是确保注意力权重的形状和大小正确以及在每个 UNet 块中设置适当数量的注意力处理器的重要步骤。
st = unet.state_dict()
for name, _ in unet.attn_processors.items():
cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
if name.startswith("mid_block"):
hidden_size = unet.config.block_out_channels[-1]
elif name.startswith("up_blocks"):
block_id = int(name[len("up_blocks.")])
hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
elif name.startswith("down_blocks"):
block_id = int(name[len("down_blocks.")])
hidden_size = unet.config.block_out_channels[block_id]
layer_name = name.split(".processor")[0]
weights = {
"to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
"to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
}
if train_q_out:
weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
if cross_attention_dim is not None:
custom_diffusion_attn_procs[name] = attention_class(
train_kv=train_kv,
train_q_out=train_q_out,
hidden_size=hidden_size,
cross_attention_dim=cross_attention_dim,
).to(unet.device)
custom_diffusion_attn_procs[name].load_state_dict(weights)
else:
custom_diffusion_attn_procs[name] = attention_class(
train_kv=False,
train_q_out=False,
hidden_size=hidden_size,
cross_attention_dim=cross_attention_dim,
)
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)
优化器 被初始化以更新交叉注意力层的参数:
optimizer = optimizer_class(
itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
if args.modifier_token is not None
else custom_diffusion_layers.parameters(),
lr=args.learning_rate,
betas=(args.adam_beta1, args.adam_beta2),
weight_decay=args.adam_weight_decay,
eps=args.adam_epsilon,
)
在训练循环中,重要的是只更新你试图学习的概念的嵌入。这意味着将所有其他标记嵌入的梯度设置为零:
if args.modifier_token is not None:
if accelerator.num_processes > 1:
grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
else:
grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
for i in range(len(modifier_token_id[1:])):
index_grads_to_zero = index_grads_to_zero & (
torch.arange(len(tokenizer)) != modifier_token_id[i]
)
grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
index_grads_to_zero, :
].fill_(0)
启动脚本
一旦你完成了所有更改或对默认配置满意,就可以启动训练脚本了!🚀
在本指南中,你将下载并使用这些示例猫图片。你也可以创建并使用自己的数据集(参见为训练创建数据集指南)。
将环境变量 MODEL_NAME
设置为 Hub 上的模型 ID 或本地模型的路径,INSTANCE_DIR
设置为你刚刚下载猫图片的路径,OUTPUT_DIR
设置为你希望保存模型的位置。你将使用 `
一旦训练完成,你可以使用你的新自定义扩散模型进行推理。
下一步
恭喜你使用自定义扩散训练了一个模型!🎉 要了解更多:
- 阅读 多概念自定义文本到图像扩散 博客文章,了解自定义扩散团队的实验结果详情。