InstructPix2Pix

InstructPix2Pix — это модель Stable Diffusion, обученная редактировать изображения по инструкциям, заданным человеком. Например, ваш промпт может звучать как «сделай облака дождливыми», и модель отредактирует входное изображение соответственно. Эта модель зависит от текстового промпта (или инструкции по редактированию) и входного изображения.

В этом руководстве мы рассмотрим скрипт обучения train_instruct_pix2pix.py, чтобы вы могли ознакомиться с ним и понять, как адаптировать его под собственные задачи.

Перед запуском скрипта убедитесь, что вы установили библиотеку из исходного кода:

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

Затем перейдите в папку с примером, содержащую скрипт обучения, и установите необходимые зависимости для используемого вами скрипта:

cd examples/instruct_pix2pix
pip install -r requirements.txt

🤗 Accelerate — это библиотека, которая помогает обучать модели на нескольких GPU/TPU или с использованием смешанной точности. Она автоматически настраивает среду обучения в зависимости от вашего оборудования и окружения. Подробнее можно узнать в Кратком обзоре 🤗 Accelerate.

Инициализируйте среду 🤗 Accelerate:

accelerate config

Чтобы настроить среду по умолчанию 🤗 Accelerate без выбора каких-либо конфигураций:

accelerate config default

Или, если ваша среда не поддерживает интерактивную оболочку, например, ноутбук, вы можете использовать:

from accelerate.utils import write_basic_config

write_basic_config()

Наконец, если вы хотите обучить модель на собственном наборе данных, ознакомьтесь с руководством Create a dataset for training, чтобы узнать, как создать набор данных, совместимый со скриптом обучения.

В следующих разделах выделены части скрипта обучения, важные для понимания того, как его модифицировать, но не все аспекты скрипта описаны подробно. Если вы хотите узнать больше, не стесняйтесь прочитать скрипт и сообщить нам, если у вас возникнут вопросы или замечания.

Параметры скрипта

Скрипт обучения содержит множество параметров, позволяющих настроить процесс обучения. Все параметры и их описания находятся в функции parse_args(). Для большинства параметров заданы значения по умолчанию, которые хорошо работают, но при необходимости вы можете указать свои значения в команде запуска обучения.

Например, чтобы увеличить разрешение входного изображения:

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

Многие базовые и важные параметры описаны в руководстве по обучению Text-to-image, поэтому в данном руководстве рассматриваются только параметры, относящиеся к InstructPix2Pix:

• --original_image_column: исходное изображение до внесения правок
• --edited_image_column: изображение после внесения правок
• --edit_prompt_column: инструкции для редактирования изображения
• --conditioning_dropout_prob: вероятность dropout для отредактированного изображения и промптов редактирования во время обучения, которая позволяет использовать classifier-free guidance (CFG) для одного или обоих условных входов

Скрипт обучения

Код предобработки набора данных и цикл обучения находятся в функции main(). Именно здесь вы вносите изменения в скрипт обучения, чтобы адаптировать его под свою задачу.

Как и в случае с параметрами скрипта, пошаговое описание скрипта обучения приведено в руководстве по обучению Text-to-image. В данном руководстве рассматриваются только те части скрипта, которые относятся к InstructPix2Pix.

Скрипт начинается с изменения количества входных каналов в первом свёрточном слое UNet, чтобы учесть дополнительное условное изображение InstructPix2Pix:

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 обновляются оптимизатором:

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,
)

Далее отредактированные изображения и инструкции по редактированию предварительно обрабатываются и токенизируются. Важно, чтобы к исходным и отредактированным изображениям применялись одни и те же преобразования.

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

Наконец, в цикле обучения сначала кодируются отредактированные изображения в латентное пространство:

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

Затем скрипт применяет dropout к эмбеддингам исходного изображения и инструкции по редактированию для поддержки CFG. Именно это позволяет модели модулировать влияние инструкции по редактированию и исходного изображения на отредактированное изображение.

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

That’s pretty much it! Aside from the differences described here, the rest of the script is very similar to the Text-to-image training script, so feel free to check it out for more details. If you want to learn more about how the training loop works, check out the Understanding pipelines, models and schedulers tutorial which breaks down the basic pattern of the denoising process.

Запустите скрипт

Когда вы будете довольны внесёнными в скрипт изменениями или сойдётесь с конфигурацией по умолчанию, можно запускать скрипт обучения! 🚀

В этом руководстве используется датасет fusing/instructpix2pix-1000-samples, представляющий собой уменьшенную версию исходного датасета. При желании вы можете создать и использовать собственный датасет (см. руководство Create a dataset for training).

Установите переменную окружения MODEL_NAME в имя модели (это может быть идентификатор модели на Hub или путь к локальной модели), а DATASET_ID — в имя набора данных на Hub. Скрипт создаёт и сохраняет все компоненты (feature extractor, scheduler, text encoder, UNet и т. д.) в подпапку вашего репозитория.

Для лучших результатов попробуйте более длительные циклы обучения с использованием набора данных большего размера. Мы тестировали этот скрипт обучения только на наборе данных меньшего масштаба.

Для отслеживания хода обучения с помощью Weights and Biases добавьте параметр --report_to=wandb к команде обучения, а также укажите изображение для валидации с помощью --val_image_url и промпт для валидации с помощью --validation_prompt. Это может быть очень полезно для отладки модели.

Если вы обучаете модель на нескольких GPU, добавьте параметр --multi_gpu к команде accelerate launch.

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 для инференса:

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

pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained("your_cool_model", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
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_steps, image_guidance_scale и guidance_scale, чтобы оценить их влияние на скорость и качество вывода. Параметры масштаба руководства особенно значимы, поскольку они определяют степень влияния исходного изображения и инструкций по редактированию на итоговый результат.

Stable Diffusion XL

Stable Diffusion XL (SDXL) — это мощная модель text-to-image, генерирующая изображения высокого разрешения, которая добавляет второй text-encoder в свою архитектуру. Используйте скрипт train_instruct_pix2pix_sdxl.py для обучения модели SDXL выполнять инструкции по редактированию изображений.

Скрипт обучения SDXL подробно описан в руководстве обучение SDXL.

Дальнейшие шаги

Поздравляем с обучением собственной модели InstructPix2Pix! 🥳 Чтобы узнать о ней больше, вам может помочь:

• Ознакомьтесь с статьей в блоге «Instruction-tuning Stable Diffusion with InstructPix2Pix», чтобы узнать больше о наших экспериментах с InstructPix2Pix, подготовке датасета и результатах для различных инструкций.