Texte vers image

Le script de génération d’images à partir de texte 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.

Les modèles de texte vers image tels que Stable Diffusion sont conditionnés pour générer des images à partir d'un prompt texte.

L’entraînement d’un modèle peut solliciter fortement votre matériel, mais si vous activez gradient_checkpointing et mixed_precision, il est possible d’entraîner un modèle sur un seul GPU de 24 Go. Si vous entraînez avec des tailles de lot plus importantes ou souhaitez accélérer l’entraînement, il est préférable d’utiliser des GPU dotés de plus de 30 Go de mémoire. Vous pouvez réduire votre empreinte mémoire en activant l’attention efficace en mémoire avec xFormers.

Ce guide explore le script d’entraînement train_text_to_image.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.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 fp16, ajoutez le paramètre --mixed_precision à la commande d’entraînement :

accelerate launch train_text_to_image.py \
  --mixed_precision="fp16"

Certains paramètres de base et importants incluent :

• --pretrained_model_name_or_path : le nom du modèle sur le Hub ou un chemin local vers le modèle pré-entraîné
• --dataset_name : le nom du jeu de données sur le Hub ou un chemin local vers le jeu de données sur lequel entraîner le modèle
• --image_column : le nom de la colonne d'images dans le jeu de données sur lequel entraîner le modèle
• --caption_column : le nom de la colonne de texte du jeu de données sur lequel entraîner
• --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_text_to_image.py \
  --snr_gamma=5.0

Vous pouvez comparer les surfaces de perte pour différentes valeurs de snr_gamma dans ce rapport Weights and Biases. Pour les jeux de données plus petits, les effets de Min-SNR peuvent ne pas être aussi évidents que pour les jeux de données plus volumineux.

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(). Si vous devez adapter le script d’entraînement, c’est ici que vous devrez apporter vos modifications.

Le script train_text_to_image commence par charger un scheduler et un tokenizer. Vous pouvez choisir d'utiliser un scheduler différent ici si vous le souhaitez :

noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
)

Ensuite, le script charge le modèle UNet :

load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
model.register_to_config(**load_model.config)

model.load_state_dict(load_model.state_dict())

Ensuite, les colonnes texte et image du jeu de données doivent être prétraitées. La fonction tokenize_captions gère la tokenisation des entrées, et la fonction train_transforms spécifie le type de transformations à appliquer à l'image. Ces deux fonctions sont regroupées dans preprocess_train :

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["pixel_values"] = [train_transforms(image) for image in images]
    examples["input_ids"] = tokenize_captions(examples)
    return examples

Enfin, la boucle d’entraînement gère tout le reste. Elle encode les images dans l’espace latent, ajoute du bruit aux latents, calcule les embeddings de texte sur lesquels conditionner, met à jour les paramètres du modèle, puis enregistre et pousse le modèle sur le Hub. 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înez-vous 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). Si vous entraînez sur plus d’une GPU, ajoutez le paramètre --multi_gpu à la commande accelerate launch.

Pour entraîner sur un jeu de données local, définissez les variables d'environnement TRAIN_DIR et OUTPUT_DIR sur le chemin du jeu de données et l'emplacement où enregistrer le modèle.

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export dataset_name="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --enable_xformers_memory_efficient_attention \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model" \
  --push_to_hub

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

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("path/to/saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

image = pipeline(prompt="yoda").images[0]
image.save("yoda-naruto.png")

Étapes suivantes

Félicitations pour avoir entraîné votre propre modèle de texte vers image ! Pour en savoir plus sur l'utilisation de votre nouveau modèle, les guides suivants peuvent vous être utiles :

• Apprenez à charger les poids LoRA pour l'inférence si vous avez entraîné votre modèle avec LoRA.
• En savoir plus sur la manière dont certains paramètres comme l’échelle de guidage ou des techniques telles que le pondération des prompts peuvent vous aider à contrôler l’inférence dans le guide de la tâche Text-to-image.