InstructPix2Pix

InstructPix2Pix é um modelo Stable Diffusion treinado para editar imagens a partir de instruções fornecidas por humanos. Por exemplo, seu prompt pode ser “turn the clouds rainy” e o modelo editará a imagem de entrada de acordo. Este modelo é condicionado no prompt de texto (ou instrução de edição) e na imagem de entrada.

Este guia explorará o script de treinamento train_instruct_pix2pix.py para ajudá-lo a se familiarizar com ele e como você pode adaptá-lo para o seu próprio caso de uso.

Antes de executar o script, certifique-se de instalar a biblioteca a partir do código-fonte:

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

Em seguida, navegue até a pasta de exemplo que contém o script de treinamento e instale as dependências necessárias para o script que você está usando:

cd examples/instruct_pix2pix
pip install -r requirements.txt

🤗 Accelerate é uma biblioteca para ajudar você a treinar em múltiplos GPUs/TPUs ou com precisão mista. Ela configura automaticamente seu ambiente de treinamento com base no seu hardware e ambiente. Dê uma olhada no Quick tour do 🤗 Accelerate para saber mais.

Inicialize um ambiente do 🤗 Accelerate:

accelerate config

Para configurar um ambiente padrão do 🤗 Accelerate sem escolher nenhuma configuração:

accelerate config default

Ou, se o seu ambiente não suporta um shell interativo, como um notebook, você pode usar:

from accelerate.utils import write_basic_config

write_basic_config()

Por fim, se você deseja treinar um modelo em seu próprio conjunto de dados, consulte o guia Criar um conjunto de dados para treinamento para aprender a criar um conjunto de dados que funcione com o script de treinamento.

As seções a seguir destacam partes do script de treinamento que são importantes para entender como modificá-lo, mas não cobrem todos os aspectos do script em detalhes. Se você tiver interesse em aprender mais, sinta-se à vontade para ler o script e nos informar caso tenha alguma dúvida ou preocupação.

Parâmetros do script

O script de treinamento possui muitos parâmetros para ajudar você a personalizar sua execução de treinamento. Todos os parâmetros e suas descrições estão na função parse_args(). Valores padrão são fornecidos para a maioria dos parâmetros e funcionam muito bem, mas você também pode definir seus próprios valores no comando de treinamento, se desejar.

Por exemplo, para aumentar a resolução da imagem de entrada:

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

Muitos dos parâmetros básicos e importantes são descritos no guia de treinamento Texto-para-imagem, portanto, este guia se concentra apenas nos parâmetros relevantes para o InstructPix2Pix:

• --original_image_column: a imagem original antes das edições serem feitas
• --edited_image_column: a imagem após a aplicação das edições
• --edit_prompt_column: as instruções para editar a imagem
• --conditioning_dropout_prob: a probabilidade de dropout para a imagem editada e os prompts de edição durante o treinamento, que habilita o classifier-free guidance (CFG) para um ou ambos os inputs de condicionamento

Script de treinamento

O código de pré-processamento do conjunto de dados e o loop de treinamento estão na função main(). É aqui que você fará suas alterações no script de treinamento para adaptá-lo ao seu próprio caso de uso.

Assim como com os parâmetros do script, um passo a passo do script de treinamento é fornecido no guia de treinamento Text-to-image. Em vez disso, este guia examina as partes do script relevantes para o InstructPix2Pix.

O script começa modificando o número de canais de entrada na primeira camada convolucional do UNet para levar em conta a imagem de condicionamento adicional do 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

Esses parâmetros do UNet são atualizados pelo otimizador:

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

Em seguida, as imagens editadas e as instruções de edição são pré-processadas e tokenizadas. É importante que as mesmas transformações de imagem sejam aplicadas às imagens originais e editadas.

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

Por fim, no loop de treinamento, ele começa codificando as imagens editadas no espaço latente:

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

Em seguida, o script aplica dropout aos embeddings da imagem original e da instrução de edição para suportar o CFG. É isso que permite ao modelo modular a influência da instrução de edição e da imagem original sobre a imagem editada.

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

É praticamente isso! Além das diferenças descritas aqui, o restante do script é muito semelhante ao script de treinamento Text-to-image, então sinta-se à vontade para consultá-lo para mais detalhes. Se você quiser aprender mais sobre como o loop de treinamento funciona, confira o tutorial Understanding pipelines, models and schedulers, que detalha o padrão básico do processo de remoção de ruído.

Inicie o script

Quando você estiver satisfeito com as alterações no seu script ou estiver de acordo com a configuração padrão, estará pronto para iniciar o script de treinamento! 🚀

Este guia utiliza o conjunto de dados fusing/instructpix2pix-1000-samples, que é uma versão menor do conjunto de dados original. Você também pode criar e usar seu próprio conjunto de dados, se desejar (veja o guia Criar um conjunto de dados para treinamento).

Defina a variável de ambiente MODEL_NAME com o nome do modelo (pode ser um id de modelo no Hub ou um caminho para um modelo local) e a DATASET_ID com o nome do dataset no Hub. O script cria e salva todos os componentes (feature extractor, scheduler, text encoder, UNet, etc.) em uma subpasta do seu repositório.

Para melhores resultados, tente treinos mais longos com um conjunto de dados maior. Testamos este script de treinamento apenas em um conjunto de dados de menor escala.

Para monitorar o progresso do treinamento com Weights and Biases, adicione o parâmetro --report_to=wandb ao comando de treinamento e especifique uma imagem de validação com --val_image_url e um prompt de validação com --validation_prompt. Isso pode ser muito útil para depurar o modelo.

Se você estiver treinando em mais de uma GPU, adicione o parâmetro --multi_gpu ao comando 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

Após o treinamento ser concluído, você pode usar seu novo InstructPix2Pix para inferência:

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

Você deve experimentar diferentes valores de num_inference_steps, image_guidance_scale e guidance_scale para ver como eles afetam a velocidade e a qualidade da inferência. Os parâmetros de escala de orientação são especialmente impactantes, pois controlam o quanto a imagem original e as instruções de edição afetam a imagem editada.

Stable Diffusion XL

Stable Diffusion XL (SDXL) é um modelo poderoso de texto para imagem que gera imagens em alta resolução e adiciona um segundo text-encoder à sua arquitetura. Use o script train_instruct_pix2pix_sdxl.py para treinar um modelo SDXL a seguir instruções de edição de imagem.

O script de treinamento do SDXL é discutido em mais detalhes no guia de treinamento do SDXL.

Próximas etapas

Parabéns por treinar seu próprio modelo InstructPix2Pix! 🥳 Para saber mais sobre o modelo, pode ser útil:

• Leia o post no blog Instruction-tuning Stable Diffusion with InstructPix2Pix para saber mais sobre alguns experimentos que realizamos com InstructPix2Pix, a preparação do conjunto de dados e os resultados para diferentes instruções.