Skip to content

InstructPix2Pix

InstructPix2Pix 是一个 Stable Diffusion 模型,经过训练可以根据人类提供的指令编辑图像。例如,你的提示可以是“让云变成雨天”,模型将根据输入图像进行相应的编辑。该模型以文本提示(或编辑指令)和输入图像为条件。

本指南将探讨 train_instruct_pix2pix.py 训练脚本,帮助你熟悉它,以及如何根据你的用例进行调整。

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

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

然后导航到包含训练脚本的示例文件夹,并安装你使用的脚本所需的所有依赖项:

bash
cd examples/instruct_pix2pix
pip install -r requirements.txt

初始化一个 🤗 Accelerate 环境:

bash
accelerate config

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

bash
accelerate config default

或者如果你的环境不支持交互式 shell,比如在笔记本中,你可以使用:

py
from accelerate.utils import write_basic_config

write_basic_config()

最后,如果你想在自己的数据集上训练模型,请查看创建训练数据集指南,学习如何创建与训练脚本兼容的数据集。

脚本参数

训练脚本有许多参数,帮助你自定义训练运行。所有参数及其描述都可以在parse_args()函数中找到。大多数参数都有默认值,效果相当不错,但你也可以在训练命令中设置自己的值。

例如,要增加输入图像的分辨率:

bash
accelerate launch train_instruct_pix2pix.py \
  --resolution=512 \

许多基本和重要的参数在文本到图像训练指南中有详细描述,因此本指南仅关注与InstructPix2Pix相关的参数:

  • --original_image_column:编辑前的原始图像
  • --edited_image_column:编辑后的图像
  • --edit_prompt_column:编辑图像的指令
  • --conditioning_dropout_prob:训练期间编辑图像和编辑指令的丢弃概率,这使得无分类器指导(CFG)可以应用于一个或两个条件输入

训练脚本

数据集预处理代码和训练循环可以在main()函数中找到。这是你将修改训练脚本以适应自己用例的地方。

与脚本参数一样,训练脚本的详细介绍在文本到图像训练指南中提供。相反,本指南将查看脚本中与InstructPix2Pix相关的部分。

脚本首先修改UNet第一卷积层中的输入通道数量,以适应InstructPix2Pix的额外条件图像:

py
in_channels = 8
out_channels = unet.conv_in.out_channels
unet.register_to_config(in_channels=in_channels)

with torch.no_grad():
    new_conv_in = nn.Conv2d(
        in_channels, out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding
    )
    new_conv_in.weight.zero_()
    new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
    unet.conv_in = new_conv_in

这些UNet参数由优化器更新

py
optimizer = optimizer_cls(
    unet.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

接下来,编辑后的图像和编辑指令会进行预处理标记化。重要的是,相同的图像变换应应用于原始图像和编辑后的图像。

py
def preprocess_train(examples):
    preprocessed_images = preprocess_images(examples)

    original_images, edited_images = preprocessed_images.chunk(2)
    original_images = original_images.reshape(-1, 3, args.resolution, args.resolution)
    edited_images = edited_images.reshape(-1, 3, args.resolution, args.resolution)

    examples["original_pixel_values"] = original_images
    examples["edited_pixel_values"] = edited_images

    captions = list(examples[edit_prompt_column])
    examples["input_ids"] = tokenize_captions(captions)
    return examples

最后,在训练循环中,首先将编辑后的图像编码到潜在空间:

py
latents = vae.encode(batch["edited_pixel_values"].to(weight_dtype)).latent_dist.sample()
latents = latents * vae.config.scaling_factor

然后,脚本对原始图像和编辑指令嵌入应用了dropout,以支持CFG。这使得模型能够调节编辑指令和原始图像对编辑后图像的影响。

py
encoder_hidden_states = text_encoder(batch["input_ids"])[0]
original_image_embeds = vae.encode(batch["original_pixel_values"].to(weight_dtype)).latent_dist.mode()

if args.conditioning_dropout_prob is not None:
    random_p = torch.rand(bsz, device=latents.device, generator=generator)
    prompt_mask = random_p < 2 * args.conditioning_dropout_prob
    prompt_mask = prompt_mask.reshape(bsz, 1, 1)
    null_conditioning = text_encoder(tokenize_captions([""]).to(accelerator.device))[0]
    encoder_hidden_states = torch.where(prompt_mask, null_conditioning, encoder_hidden_states)

    image_mask_dtype = original_image_embeds.dtype
    image_mask = 1 - (
        (random_p >= args.conditioning_dropout_prob).to(image_mask_dtype)
        * (random_p < 3 * args.conditioning_dropout_prob).to(image_mask_dtype)
    )
    image_mask = image_mask.reshape(bsz, 1, 1, 1)
    original_image_embeds = image_mask * original_image_embeds

基本上就是这样了!除了这里描述的差异之外,脚本的其余部分与文本到图像训练脚本非常相似,因此欢迎你查看更多详细信息。如果你想了解更多关于训练循环的工作原理,请查看理解管道、模型和调度器教程,该教程分解了去噪过程的基本模式。

启动脚本

一旦你对脚本的更改感到满意,或者如果你对默认配置没问题,就可以准备启动训练脚本了!🚀

本指南使用fusing/instructpix2pix-1000-samples数据集,这是原始数据集的一个较小版本。如果你愿意,也可以创建并使用自己的数据集(参见创建训练数据集指南)。

MODEL_NAME环境变量设置为模型名称(可以是Hub上的模型ID或本地模型的路径),并将DATASET_ID设置为Hub上数据集的名称。脚本会在你的仓库的子文件夹中创建并保存所有组件(特征提取器、调度器、文本编码器、UNet等)。

如果你在多个GPU上进行训练,请在accelerate launch命令中添加--multi_gpu参数。

bash
accelerate launch --mixed_precision="fp16" train_instruct_pix2pix.py \
    --pretrained_model_name_or_path=$MODEL_NAME \
    --dataset_name=$DATASET_ID \
    --enable_xformers_memory_efficient_attention \
    --resolution=256 \
    --random_flip \
    --train_batch_size=4 \
    --gradient_accumulation_steps=4 \
    --gradient_checkpointing \
    --max_train_steps=15000 \
    --checkpointing_steps=5000 \
    --checkpoints_total_limit=1 \
    --learning_rate=5e-05 \
    --max_grad_norm=1 \
    --lr_warmup_steps=0 \
    --conditioning_dropout_prob=0.05 \
    --mixed_precision=fp16 \
    --seed=42 \
    --push_to_hub

训练完成后,你可以使用新的 InstructPix2Pix 进行推理:

py
import PIL
import requests
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline
from diffusers.utils import load_image

pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained("your_cool_model", torch_dtype=torch.float16).to("cuda")
generator = torch.Generator("cuda").manual_seed(0)

image = load_image("https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/test_pix2pix_4.png")
prompt = "add some ducks to the lake"
num_inference_steps = 20
image_guidance_scale = 1.5
guidance_scale = 10

edited_image = pipeline(
   prompt,
   image=image,
   num_inference_steps=num_inference_steps,
   image_guidance_scale=image_guidance_scale,
   guidance_scale=guidance_scale,
   generator=generator,
).images[0]
edited_image.save("edited_image.png")

你应该尝试不同的 num_inference_stepsimage_guidance_scaleguidance_scale 值,以了解它们如何影响推理速度和质量。指导尺度参数尤其具有影响力,因为它们控制原始图像和编辑指令对编辑后图像的影响程度。

Stable Diffusion XL

Stable Diffusion XL (SDXL) 是一个强大的文本到图像模型,能够生成高分辨率图像,并在其架构中添加了第二个文本编码器。使用 train_instruct_pix2pix_sdxl.py 脚本来训练一个 SDXL 模型,使其遵循图像编辑指令。

SDXL 训练脚本在 SDXL 训练 指南中有更详细的讨论。

下一步

恭喜你训练了自己的 InstructPix2Pix 模型!🥳 要了解更多关于该模型的信息,可以参考以下内容: