Stable Diffusion XL

Ce script est expérimental, et il est facile de sur-apprendre et de rencontrer des problèmes tels que l'oubli catastrophique. Essayez d'explorer différents hyperparamètres pour obtenir les meilleurs résultats sur votre jeu de données.

Stable Diffusion XL (SDXL) est une itération plus grande et plus puissante du modèle Stable Diffusion, capable de produire des images en plus haute résolution.

L’UNet de SDXL est trois fois plus volumineux et le modèle ajoute un second encodeur de texte à l’architecture. Selon le matériel dont vous disposez, cela peut être très intensif en calculs et le modèle peut ne pas s’exécuter sur un GPU grand public comme une Tesla T4. Pour aider à faire tenir ce modèle plus volumineux en mémoire et accélérer l’entraînement, essayez d’activer gradient_checkpointing, mixed_precision et gradient_accumulation_steps. Vous pouvez réduire encore davantage votre consommation de mémoire en activant l’attention efficace en mémoire avec xFormers et en utilisant l’optimiseur 8 bits de bitsandbytes.

Ce guide explore le script d’entraînement train_text_to_image_sdxl.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 .

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/text_to_image
pip install -r requirements_sdxl.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.

Paramètres du script

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.

Le script d’entraînement fournit 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(). Cette fonction fournit des valeurs par défaut pour chaque paramètre, telles que la taille du lot d’entraînement et le taux d’apprentissage, mais vous pouvez également définir vos propres valeurs dans la commande d’entraînement si vous le souhaitez.

Par exemple, pour accélérer l’entraînement avec la précision mixte en utilisant le format bf16, ajoutez le paramètre --mixed_precision à la commande d’entraînement :

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

La plupart des paramètres sont identiques à ceux du guide d’entraînement Text-to-image, vous vous concentrerez donc sur les paramètres pertinents pour l’entraînement de SDXL dans ce guide.

• --pretrained_vae_model_name_or_path : chemin vers un VAE pré-entraîné ; le VAE de SDXL est connu pour souffrir d'instabilité numérique, ce paramètre vous permet donc de spécifier un VAE plus performant
• --proportion_empty_prompts : la proportion de prompts d'image à remplacer par des chaînes vides
• --timestep_bias_strategy : indique à quel moment (début ou fin) de la séquence de timesteps appliquer un biais, ce qui peut inciter le modèle à apprendre des détails de basse ou de haute fréquence
• --timestep_bias_multiplier : le poids du biais à appliquer au timestep
• --timestep_bias_begin : le timestep à partir duquel appliquer le biais
• --timestep_bias_end : le timestep à partir duquel l'application du biais se termine
• --timestep_bias_portion : la proportion de timesteps à laquelle appliquer le biais

Pondération Min-SNR

La stratégie de pondération Min-SNR peut faciliter l’entraînement en rééquilibrant la perte pour obtenir 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_text_to_image_sdxl.py \
  --snr_gamma=5.0

Script d'entraînement

Le script d’entraînement est également similaire au guide d’entraînement Text-to-image, mais il a été modifié pour prendre en charge l’entraînement de SDXL. Ce guide se concentrera sur le code propre au script d’entraînement de SDXL.

On commence par créer des fonctions pour tokeniser les prompts afin de calculer les embeddings de prompts, et pour calculer les embeddings d’images avec le VAE. Ensuite, vous créerez une fonction pour générer les poids des timesteps en fonction du nombre de timesteps et de la stratégie de biais de timestep à appliquer.

Dans la fonction main(), en plus du chargement d'un tokenizer, le script charge un second tokenizer et un encodeur de texte, car l'architecture SDXL en utilise deux de chaque :

tokenizer_one = AutoTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision, use_fast=False
)
tokenizer_two = AutoTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer_2", revision=args.revision, use_fast=False
)

text_encoder_cls_one = import_model_class_from_model_name_or_path(
    args.pretrained_model_name_or_path, args.revision
)
text_encoder_cls_two = import_model_class_from_model_name_or_path(
    args.pretrained_model_name_or_path, args.revision, subfolder="text_encoder_2"
)

Les embeddings de prompt et d’image sont d’abord calculés et conservés en mémoire, ce qui ne pose généralement pas de problème pour un jeu de données de petite taille, mais peut entraîner des problèmes de mémoire pour des jeux de données plus volumineux. Dans ce cas, vous devez enregistrer les embeddings précalculés sur le disque séparément et les charger en mémoire pendant le processus d’entraînement (voir cette PR pour une discussion plus approfondie sur ce sujet).

text_encoders = [text_encoder_one, text_encoder_two]
tokenizers = [tokenizer_one, tokenizer_two]
compute_embeddings_fn = functools.partial(
    encode_prompt,
    text_encoders=text_encoders,
    tokenizers=tokenizers,
    proportion_empty_prompts=args.proportion_empty_prompts,
    caption_column=args.caption_column,
)

train_dataset = train_dataset.map(compute_embeddings_fn, batched=True, new_fingerprint=new_fingerprint)
train_dataset = train_dataset.map(
    compute_vae_encodings_fn,
    batched=True,
    batch_size=args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps,
    new_fingerprint=new_fingerprint_for_vae,
)

Après le calcul des embeddings, l'encodeur de texte, le VAE et le tokenizer sont supprimés pour libérer de la mémoire :

del text_encoders, tokenizers, vae
gc.collect()
torch.cuda.empty_cache()

Enfin, la boucle d’entraînement s’occupe du reste. Si vous avez choisi d’appliquer une stratégie de biais sur les pas de temps, vous verrez que les poids des pas de temps sont calculés et ajoutés en tant que bruit :

weights = generate_timestep_weights(args, noise_scheduler.config.num_train_timesteps).to(
        model_input.device
    )
    timesteps = torch.multinomial(weights, bsz, replacement=True).long()

noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)

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

Une fois que vous avez apporté toutes vos modifications ou que vous êtes satisfait de la configuration par défaut, vous êtes prêt à lancer le script d’entraînement ! 🚀

Entraînons-nous sur le jeu de données Naruto BLIP captions pour générer vos propres personnages de Naruto. Définissez les variables d’environnement MODEL_NAME et DATASET_NAME sur le modèle et le jeu de données (soit depuis le Hub, soit un chemin local). Vous devez également spécifier un VAE autre que le VAE de SDXL (soit depuis le Hub, soit un chemin local) avec VAE_NAME afin d’éviter les instabilités numériques.

Pour suivre la progression de l’entraînement avec Weights & Biases, ajoutez le paramètre --report_to=wandb à la commande d’entraînement. Vous devrez également ajouter --validation_prompt et --validation_epochs à la commande d’entraînement pour suivre les résultats. Cela peut s’avérer très utile pour déboguer le modèle et consulter les résultats intermédiaires.

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fix"
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch train_text_to_image_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --pretrained_vae_model_name_or_path=$VAE_NAME \
  --dataset_name=$DATASET_NAME \
  --enable_xformers_memory_efficient_attention \
  --resolution=512 \
  --center_crop \
  --random_flip \
  --proportion_empty_prompts=0.2 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=10000 \
  --use_8bit_adam \
  --learning_rate=1e-06 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --mixed_precision="fp16" \
  --report_to="wandb" \
  --validation_prompt="a cute Sundar Pichai creature" \
  --validation_epochs 5 \
  --checkpointing_steps=5000 \
  --output_dir="sdxl-naruto-model" \
  --push_to_hub

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

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path/to/your/model", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A naruto with green eyes and red legs."
image = pipeline(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("naruto.png")

from diffusers import DiffusionPipeline
import torch
import torch_xla.core.xla_model as xm

device = xm.xla_device()
pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to(device)

prompt = "A naruto with green eyes and red legs."
start = time()
image = pipeline(prompt, num_inference_steps=inference_steps).images[0]
print(f'Compilation time is {time()-start} sec')
image.save("naruto.png")

start = time()
image = pipeline(prompt, num_inference_steps=inference_steps).images[0]
print(f'Inference time is {time()-start} sec after compilation')

Étapes suivantes

Félicitations pour l'entraînement d'un modèle SDXL ! Pour en savoir plus sur l'utilisation de votre nouveau modèle, les guides suivants peuvent s'avérer utiles :

• Lisez le guide Stable Diffusion XL pour apprendre à l’utiliser pour diverses tâches (texte-vers-image, image-vers-image, inpainting), à utiliser son modèle de raffinement, ainsi que les différents types de micro-conditionnements.
• Consultez les guides d’entraînement DreamBooth et LoRA pour apprendre à entraîner un modèle SDXL personnalisé à partir de quelques images d’exemple seulement. Ces deux techniques d’entraînement peuvent même être combinées !