DreamBooth

DreamBooth est une technique d’entraînement qui met à jour le modèle de diffusion complet en s’entraînant sur seulement quelques images d’un sujet ou d’un style. Elle fonctionne en associant un mot spécial dans le prompt aux images d’exemple.

Pour charger un point de contrôle entraîné pour l'inférence, consultez Charger un modèle DreamBooth pour l'inférence.

Si vous entraînez sur un GPU doté d’une vRAM limitée, vous devriez essayer d’activer les paramètres gradient_checkpointing et mixed_precision dans la commande d’entraînement. Vous pouvez également réduire votre empreinte mémoire en utilisant une attention efficace en mémoire avec xFormers.

Ce guide examine le script train_dreambooth.py afin de vous familiariser avec celui-ci et de vous montrer 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 .

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/dreambooth
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

DreamBooth est très sensible aux hyperparamètres d’entraînement et il est facile de sur-apprendre. Lisez le billet de blog Training Stable Diffusion with Dreambooth using 🧨 Diffusers pour obtenir les paramètres recommandés pour différents sujets afin de vous aider à choisir les hyperparamètres appropriés.

Le script d’entraînement propose de nombreux paramètres pour personnaliser votre session d’entraînement. Tous les paramètres et leurs descriptions se trouvent dans la fonction parse_args(). Les paramètres sont définis avec des valeurs par défaut qui devraient fonctionner très bien sans configuration supplémentaire, mais vous pouvez également définir vos propres valeurs dans la commande d’entraînement si vous le souhaitez.

Par exemple, pour entraîner au format bf16 :

accelerate launch train_dreambooth.py \
    --mixed_precision="bf16"

Voici quelques paramètres de base et importants à connaître et à spécifier :

--pretrained_model_name_or_path : le nom du modèle sur le Hub ou un chemin local vers le modèle pré-entraîné
--instance_data_dir : chemin vers un dossier contenant le jeu de données d'entraînement (images d'exemple)
--instance_prompt : le prompt de texte contenant le mot spécial pour les images d'exemple
--train_text_encoder : indique s'il faut entraîner également le text encoder
--output_dir : l’endroit où enregistrer le modèle entraîné
--push_to_hub : indique s'il faut pousser le modèle entraîné sur le Hub
--checkpointing_steps : fréquence de sauvegarde d'un point de contrôle pendant l'entraînement du modèle ; cela est utile si l'entraînement est interrompu pour une raison quelconque, car vous pouvez reprendre l'entraînement à partir de ce point de contrôle en ajoutant --resume_from_checkpoint à votre commande d'entraînement

Pondération Min-SNR

La stratégie de pondération Min-SNR peut aider à l'entraînement en rééquilibrant la perte pour atteindre une convergence plus rapide. Le script d'entraînement prend en charge la prédiction de epsilon (bruit) ou de v_prediction, mais Min-SNR est compatible avec les deux types de prédiction. Cette stratégie de pondération n'est prise en charge que par PyTorch.

Ajoutez le paramètre --snr_gamma et définissez-le à la valeur recommandée de 5,0 :

accelerate launch train_dreambooth.py \
  --snr_gamma=5.0

Perte de préservation des antécédents

La perte de préservation des a priori est une méthode qui utilise les échantillons générés par le modèle lui-même pour l’aider à apprendre à générer des images plus variées. Comme ces images générées appartiennent à la même classe que les images que vous avez fournies, elles aident le modèle à conserver ce qu’il a appris sur la classe et à exploiter ses connaissances existantes sur la classe pour créer de nouvelles compositions.

--with_prior_preservation : indique s'il faut utiliser la perte de préservation a priori
--prior_loss_weight : contrôle l'influence de la perte de préservation du prior sur le modèle
--class_data_dir : chemin vers un dossier contenant les images d'échantillons de classes générées
--class_prompt : le prompt de texte décrivant la classe des images générées

accelerate launch train_dreambooth.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="path/to/class/images" \
  --class_prompt="text prompt describing class"

Entraîner l'encodeur de texte

Pour améliorer la qualité des sorties générées, vous pouvez également entraîner l'encodeur de texte en plus du UNet. Cela nécessite une mémoire supplémentaire et vous aurez besoin d'un GPU doté d'au moins 24 Go de vRAM. Si vous disposez du matériel nécessaire, l'entraînement de l'encodeur de texte produit de meilleurs résultats, en particulier lors de la génération d'images de visages. Activez cette option en :

accelerate launch train_dreambooth.py \
  --train_text_encoder

Script d'entraînement

DreamBooth est livré avec ses propres classes de jeu de données :

DreamBoothDataset : prétraite les images et les images de classe, et tokenise les prompts pour l'entraînement
PromptDataset : génère les embeddings de prompt pour générer les images de classe

Si vous avez activé la perte de préservation antérieure, les images de classe sont générées ici :

sample_dataset = PromptDataset(args.class_prompt, num_new_images)
sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)

sample_dataloader = accelerator.prepare(sample_dataloader)
pipeline.to(accelerator.device)

for example in tqdm(
    sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
):
    images = pipeline(example["prompt"]).images

Ensuite, la fonction main() gère la configuration du jeu de données pour l'entraînement ainsi que la boucle d'entraînement elle-même. Le script charge le tokenizer, le scheduler et les modèles :

# Load the tokenizer
if args.tokenizer_name:
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, revision=args.revision, use_fast=False)
elif args.pretrained_model_name_or_path:
    tokenizer = AutoTokenizer.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="tokenizer",
        revision=args.revision,
        use_fast=False,
    )

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = text_encoder_cls.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)

if model_has_vae(args):
    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision
    )
else:
    vae = None

unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

Ensuite, il est temps de créer le jeu de données d'entraînement et le DataLoader à partir de DreamBoothDataset :

train_dataset = DreamBoothDataset(
    instance_data_root=args.instance_data_dir,
    instance_prompt=args.instance_prompt,
    class_data_root=args.class_data_dir if args.with_prior_preservation else None,
    class_prompt=args.class_prompt,
    class_num=args.num_class_images,
    tokenizer=tokenizer,
    size=args.resolution,
    center_crop=args.center_crop,
    encoder_hidden_states=pre_computed_encoder_hidden_states,
    class_prompt_encoder_hidden_states=pre_computed_class_prompt_encoder_hidden_states,
    tokenizer_max_length=args.tokenizer_max_length,
)

train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=args.train_batch_size,
    shuffle=True,
    collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
    num_workers=args.dataloader_num_workers,
)

Enfin, la boucle d’entraînement gère les étapes restantes, telles que la conversion des images en espace latent, l’ajout de bruit à l’entrée, la prédiction du résidu de bruit et le calcul de la perte.

Si vous souhaitez en savoir plus sur le fonctionnement de la boucle d’entraînement, consultez le tutoriel Comprendre les pipelines, les modèles et les planificateurs, qui détaille le schéma de base du processus de débruitage.

Lancer le script

Vous êtes maintenant prêt à lancer le script d'entraînement ! 🚀

Pour ce guide, vous téléchargerez quelques images de chien et les stockerez dans un répertoire. Mais rappelez-vous que vous pouvez 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).

from huggingface_hub import snapshot_download

local_dir = "./dog"
snapshot_download(
    "diffusers/dog-example",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

Définissez la variable d’environnement MODEL_NAME sur un identifiant de modèle du Hub ou un chemin vers un modèle local, INSTANCE_DIR sur le chemin où vous venez de télécharger les images de chiens, et OUTPUT_DIR sur l’endroit où vous souhaitez enregistrer le modèle. Vous utiliserez sks comme mot spécial pour lier l’entraînement.

Si vous souhaitez suivre le processus d’entraînement, vous pouvez enregistrer périodiquement les images générées au fil de l’entraînement. Ajoutez les paramètres suivants à la commande d’entraînement :

--validation_prompt="a photo of a sks dog"
--num_validation_images=4
--validation_steps=100

Une dernière chose avant de lancer le script ! Selon le GPU dont vous disposez, vous devrez peut-être activer certaines optimisations pour entraîner DreamBooth.

16 Go

12 Go

8 Go

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export INSTANCE_DIR="./dog"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=400 \
  --push_to_hub

Une fois l’entraînement terminé, vous pouvez utiliser votre nouveau modèle pour l’inférence !

Envie d’essayer votre modèle pour l’inférence avant la fin de l’entraînement ? 🤭 Assurez-vous d’avoir installé la dernière version de 🤗 Accelerate.

from diffusers import DiffusionPipeline, UNet2DConditionModel
from transformers import CLIPTextModel
import torch

unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet")

# if you have trained with `--args.train_text_encoder` make sure to also load the text encoder
text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder")

pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16,
).to("cuda")

image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path_to_saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

LoRA

LoRA est une technique d’entraînement permettant de réduire considérablement le nombre de paramètres entraînables. En conséquence, l’entraînement est plus rapide et il est plus facile de stocker les poids résultants, car ils sont beaucoup plus petits (~100 Mo). Utilisez le script train_dreambooth_lora.py pour entraîner avec LoRA.

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

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_dreambooth_lora_sdxl.py pour entraîner un modèle SDXL avec LoRA.

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

DeepFloyd IF

DeepFloyd IF est un modèle de diffusion pixel en cascade avec trois étapes. La première étape génère une image de base et les deuxième et troisième étapes augmentent progressivement la résolution de l'image de base pour obtenir une image haute résolution de 1024x1024. Utilisez les scripts train_dreambooth_lora.py ou train_dreambooth.py pour entraîner un modèle DeepFloyd IF avec LoRA ou le modèle complet.

DeepFloyd IF utilise une variance prédite, mais les scripts d’entraînement de Diffusers utilisent une erreur prédite, de sorte que les modèles DeepFloyd IF entraînés sont basculés sur un calendrier de variance fixe. Les scripts d’entraînement mettent à jour la configuration du planificateur du modèle entièrement entraîné pour vous. Toutefois, lorsque vous chargez les poids LoRA enregistrés, vous devez également mettre à jour la configuration du planificateur du pipeline.

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", use_safetensors=True)

pipe.load_lora_weights("<lora weights path>")

# Update scheduler config to fixed variance schedule
pipe.scheduler = pipe.scheduler.__class__.from_config(pipe.scheduler.config, variance_type="fixed_small")

Le modèle de l'étape 2 nécessite des images de validation supplémentaires pour l'augmentation de résolution. Vous pouvez télécharger et utiliser une version réduite des images d'entraînement à cet effet.

from huggingface_hub import snapshot_download

local_dir = "./dog_downsized"
snapshot_download(
    "diffusers/dog-example-downsized",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

Les exemples de code ci-dessous offrent un aperçu succinct de la manière d'entraîner un modèle DeepFloyd IF à l'aide d'une combinaison de DreamBooth et de LoRA. Certains paramètres importants à noter sont :

--resolution=64, une résolution beaucoup plus petite est requise car DeepFloyd IF est un modèle de diffusion de pixels et, pour travailler sur des pixels non compressés, les images d’entrée doivent être plus petites
--pre_compute_text_embeddings, calcule les embeddings de texte à l'avance pour économiser de la mémoire, car le T5Model peut consommer beaucoup de mémoire
--tokenizer_max_length=77, vous pouvez utiliser une longueur de texte par défaut plus longue avec T5 comme encodeur de texte, mais la procédure d’encodage du modèle par défaut utilise une longueur de texte plus courte
--text_encoder_use_attention_mask, pour transmettre le masque d'attention à l'encodeur de texte

Étape 1 LoRA DreamBooth

Étape 2 LoRA DreamBooth

Étape 1 DreamBooth

Étape 2 DreamBooth

L'étape 1 d'entraînement de DeepFloyd IF avec LoRA et DreamBooth nécessite environ 28 Go de mémoire.

export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth_dog_lora"

accelerate launch train_dreambooth_lora.py \
  --report_to wandb \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a sks dog" \
  --resolution=64 \
  --train_batch_size=4 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --scale_lr \
  --max_train_steps=1200 \
  --validation_prompt="a sks dog" \
  --validation_epochs=25 \
  --checkpointing_steps=100 \
  --pre_compute_text_embeddings \
  --tokenizer_max_length=77 \
  --text_encoder_use_attention_mask

Conseils d'entraînement

L’entraînement du modèle DeepFloyd IF peut s’avérer difficile, mais voici quelques conseils que nous avons trouvés utiles :

LoRA est suffisant pour l’entraînement du modèle de l’étape 1, car la faible résolution du modèle rend difficile la représentation de détails plus fins, quoi qu’il en soit.
Pour des objets courants ou simples, il n’est pas nécessaire de fine-tuner l’upscaler. Assurez-vous que le prompt transmis à l’upscaler est ajusté afin de retirer le nouveau jeton du prompt de l’instance. Par exemple, si votre prompt de l’étape 1 est « a sks dog », votre prompt de l’étape 2 doit être « a dog ».
Pour des détails plus fins, comme les visages, entraîner entièrement le redimensionneur de l’étape 2 est préférable à l’entraînement du modèle de l’étape 2 avec LoRA. Il est également utile d’utiliser des taux d’apprentissage plus faibles avec des tailles de lot plus grandes.
Des taux d'apprentissage plus faibles doivent être utilisés pour entraîner le modèle de l'étape 2.
Le DDPMScheduler fonctionne mieux que le DPMSolver utilisé dans les scripts d'entraînement.

Étapes suivantes

Félicitations pour l'entraînement de votre modèle DreamBooth ! Pour en savoir plus sur l'utilisation de votre nouveau modèle, le guide suivant peut s'avérer utile :

Apprenez à charger un point de contrôle DreamBooth pour l'inférence.
Si vous avez entraîné avec LoRA, consultez LoRA pour l'inférence.

Update on GitHub