InstructPix2Pix

InstructPix2Pix est un modèle Stable Diffusion entraîné pour modifier des images à partir d'instructions fournies par des humains. Par exemple, votre invite peut être « turn the clouds rainy » et le modèle modifiera l'image d'entrée en conséquence. Ce modèle est conditionné par l'invite texte (ou l'instruction de modification) et l'image d'entrée.

Ce guide explore le script d'entraînement train_instruct_pix2pix.py pour vous aider à le maîtriser et à savoir comment l'adapter à votre propre cas d'usage.

Avant d'exécuter le script, assurez-vous d'installer la bibliothèque à partir du code source :

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

Puis accédez au dossier d'exemple contenant le script d'entraînement et installez les dépendances requises pour le script que vous utilisez :

cd examples/instruct_pix2pix
pip install -r requirements.txt

🤗 Accelerate est une bibliothèque qui vous aide à entraîner sur plusieurs GPU/TPU ou avec une précision mixte. Elle configure automatiquement votre configuration d’entraînement en fonction de votre matériel et de votre environnement. Consultez le tour rapide de 🤗 Accelerate pour en savoir plus.

Initialiser un environnement 🤗 Accelerate :

accelerate config

Pour configurer un environnement 🤗 Accelerate par défaut sans choisir de configuration :

accelerate config default

Sinon, si votre environnement ne prend pas en charge un shell interactif, comme un notebook, vous pouvez utiliser :

from accelerate.utils import write_basic_config

write_basic_config()

Enfin, si vous souhaitez entraîner un modèle sur votre propre jeu de données, consultez le guide Create a dataset for training pour apprendre à créer un jeu de données compatible avec le script d'entraînement.

Les sections suivantes mettent en évidence les parties du script d’entraînement importantes pour comprendre comment le modifier, mais elles ne couvrent pas tous les aspects du script en détail. Si vous souhaitez en savoir plus, n’hésitez pas à lire le script et à nous faire part de vos questions ou préoccupations.

Paramètres du script

Le script d’entraînement comporte de nombreux paramètres pour vous aider à personnaliser votre session d’entraînement. Tous les paramètres et leurs descriptions se trouvent dans la fonction parse_args(). Des valeurs par défaut sont fournies pour la plupart des paramètres et fonctionnent très bien, mais vous pouvez également définir vos propres valeurs dans la commande d’entraînement si vous le souhaitez.

Par exemple, pour augmenter la résolution de l'image d'entrée :

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

De nombreux paramètres de base et importants sont décrits dans le guide d’entraînement Text-to-image, ce guide se concentre donc uniquement sur les paramètres pertinents pour InstructPix2Pix :

--original_image_column : l'image d'origine avant que les modifications ne soient effectuées
--edited_image_column : l'image après l'application des modifications
--edit_prompt_column : les instructions pour modifier l'image
--conditioning_dropout_prob : la probabilité de dropout appliquée à l'image modifiée et aux prompts de modification pendant l'entraînement, permettant le guidage sans classifieur (CFG) pour un ou les deux entrées de conditionnement

Script d'entraînement

Le code de prétraitement du jeu de données et la boucle d’entraînement se trouvent dans la fonction main(). C’est ici que vous apporterez vos modifications au script d’entraînement pour l’adapter à votre propre cas d’usage.

Comme pour les paramètres du script, un parcours détaillé du script d’entraînement est fourni dans le guide d’entraînement Text-to-image. En revanche, ce guide se concentre sur les parties du script spécifiques à InstructPix2Pix.

Le script commence par modifier le nombre de canaux d’entrée de la première couche convolutive du UNet pour prendre en compte l’image de conditionnement supplémentaire d’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

Ces paramètres UNet sont mis à jour par l’optimiseur :

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

Ensuite, les images modifiées et les instructions de modification sont prétraitées et tokenisées. Il est important que les mêmes transformations d'image soient appliquées aux images d'origine et aux images modifiées.

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

Enfin, dans la boucle d'entraînement, elle commence par encoder les images modifiées dans l'espace latent :

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

Ensuite, le script applique le dropout aux embeddings de l'image d'origine et de l'instruction d'édition pour prendre en charge le CFG. C'est ce qui permet au modèle de moduler l'influence de l'instruction d'édition et de l'image d'origine sur l'image éditée.

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

C’est à peu près tout ! À part les différences décrites ici, le reste du script est très similaire au script d’entraînement Text-to-image, n’hésitez donc pas à le consulter pour plus de détails. Si vous souhaitez en savoir plus sur le fonctionnement de la boucle d’entraînement, consultez le tutoriel Understanding pipelines, models and schedulers qui détaille le schéma de base du processus de débruitage.

Lancer le script

Une fois que vous êtes satisfait des modifications apportées à votre script ou que la configuration par défaut vous convient, vous êtes prêt à lancer le script d’entraînement ! 🚀

Ce guide utilise le jeu de données fusing/instructpix2pix-1000-samples, qui est une version réduite du jeu de données original. Vous pouvez également créer et utiliser votre propre jeu de données si vous le souhaitez (voir le guide Créer un jeu de données pour l’entraînement).

Définissez la variable d’environnement MODEL_NAME sur le nom du modèle (il peut s’agir d’un identifiant de modèle sur le Hub ou d’un chemin vers un modèle local), et la variable DATASET_ID sur le nom du jeu de données sur le Hub. Le script crée et enregistre tous les composants (extracteur de caractéristiques, planificateur, encodeur de texte, UNet, etc.) dans un sous-dossier de votre dépôt.

Pour de meilleurs résultats, essayez des sessions d’entraînement plus longues avec un jeu de données plus volumineux. Nous n’avons testé ce script d’entraînement que sur un jeu de données de plus petite envergure.

Pour suivre la progression de l’entraînement avec Weights and Biases, ajoutez le paramètre --report_to=wandb à la commande d’entraînement et spécifiez une image de validation avec --val_image_url ainsi qu’un prompt de validation avec --validation_prompt. Cela peut s’avérer très utile pour déboguer le modèle.

Si vous entraînez sur plus d’un GPU, ajoutez le paramètre --multi_gpu à la commande 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

Une fois l’entraînement terminé, vous pouvez utiliser votre nouveau InstructPix2Pix pour l’inférence :

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

Vous devriez expérimenter avec différentes valeurs de num_inference_steps, image_guidance_scale et guidance_scale pour observer leur impact sur la vitesse et la qualité de l’inférence. Les paramètres d’échelle de guidage sont particulièrement déterminants, car ils contrôlent dans quelle mesure l’image d’origine et les instructions de modification influencent l’image éditée.

Stable Diffusion XL

Stable Diffusion XL (SDXL) est un puissant modèle de génération d'images à partir de texte qui produit des images haute résolution et qui ajoute un second encodeur de texte à son architecture. Utilisez le script train_instruct_pix2pix_sdxl.py pour entraîner un modèle SDXL à suivre des instructions d'édition d'image.

Le script d’entraînement SDXL est abordé plus en détail dans le guide Entraînement SDXL.

Étapes suivantes

Félicitations pour avoir entraîné votre propre modèle InstructPix2Pix ! 🥳 Pour en savoir plus sur le modèle, il peut être utile de :

Lisez le billet de blog Instruction-tuning Stable Diffusion with InstructPix2Pix pour en savoir plus sur certaines expériences que nous avons réalisées avec InstructPix2Pix, la préparation des jeux de données et les résultats pour différentes instructions.

Update on GitHub

Diffusers

InstructPix2Pix

Paramètres du script

Script d'entraînement

Lancer le script

Stable Diffusion XL

Étapes suivantes