Inversion textuelle

Textual Inversion est une technique d’entraînement permettant de personnaliser les modèles de génération d’images à partir de quelques images d’exemple de ce que vous souhaitez qu’ils apprennent. Cette technique fonctionne en apprenant et en mettant à jour les embeddings de texte (les nouveaux embeddings sont liés à un mot spécial que vous devez utiliser dans le prompt) afin qu’ils correspondent aux images d’exemple que vous fournissez.

Pour l'inférence avec des embeddings entraînés, voir Inférence par inversion textuelle.

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 textual_inversion.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/textual_inversion
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 à adapter l’exécution d’entraînement à vos besoins. Tous les paramètres et leurs descriptions sont répertoriés dans la fonction parse_args(). Lorsque cela est applicable, Diffusers 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 n’hésitez pas à modifier ces valeurs dans la commande d’entraînement si vous le souhaitez.

Par exemple, pour augmenter le nombre d'étapes d'accumulation des gradients au-delà de la valeur par défaut de 1 :

accelerate launch textual_inversion.py \
  --gradient_accumulation_steps=4

Parmi les autres paramètres de base et importants à spécifier, on trouve :

--pretrained_model_name_or_path : le nom du modèle sur le Hub ou un chemin local vers le modèle pré-entraîné
--train_data_dir : chemin vers un dossier contenant le jeu de données d'entraînement (images d'exemple)
--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
--num_vectors : le nombre de vecteurs utilisés pour apprendre les embeddings ; augmenter ce paramètre aide le modèle à mieux apprendre, mais cela entraîne des coûts de formation plus élevés
--placeholder_token : le mot spécial à associer aux embeddings appris (vous devez utiliser ce mot dans votre prompt pour l'inférence)
--initializer_token : un mot unique qui décrit approximativement l'objet ou le style sur lequel vous essayez d'entraîner le modèle
--learnable_property : indiquez si vous entraînez le modèle pour qu’il apprenne un nouveau « style » (par exemple, le style de peinture de Van Gogh) ou un « objet » (par exemple, votre chien)

Script d'entraînement

Contrairement à certains autres scripts d’entraînement, textual_inversion.py possède une classe de jeu de données personnalisée, TextualInversionDataset, pour créer un jeu de données. Vous pouvez personnaliser la taille des images, le jeton de substitution, la méthode d’interpolation, l’application d’un recadrage, et bien plus encore. Si vous devez modifier la façon dont le jeu de données est créé, vous pouvez modifier TextualInversionDataset.

Ensuite, vous trouverez le code de prétraitement du jeu de données et la boucle d’entraînement dans la fonction main().

Le script commence par charger le tokenizer, le scheduler et le modèle :

# Load tokenizer
if args.tokenizer_name:
    tokenizer = CLIPTokenizer.from_pretrained(args.tokenizer_name)
elif args.pretrained_model_name_or_path:
    tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer")

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = CLIPTextModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

Le jeton de substitution spécial est ajouté à côté du tokenizer, et l'embedding est réajusté pour prendre en compte le nouveau jeton.

Ensuite, le script crée un jeu de données à partir de TextualInversionDataset :

train_dataset = TextualInversionDataset(
    data_root=args.train_data_dir,
    tokenizer=tokenizer,
    size=args.resolution,
    placeholder_token=(" ".join(tokenizer.convert_ids_to_tokens(placeholder_token_ids))),
    repeats=args.repeats,
    learnable_property=args.learnable_property,
    center_crop=args.center_crop,
    set="train",
)
train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.dataloader_num_workers
)

Enfin, la boucle d'entraînement gère tout le reste, de la prédiction du résidu bruité à la mise à jour des poids d'intégration du jeton de remplacement spécial.

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 ! 🚀

Pour ce guide, vous téléchargerez quelques images d'un jouet pour chat 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 = "./cat"
snapshot_download(
    "diffusers/cat_toy_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, et DATA_DIR sur le chemin où vous venez de télécharger les images de chats. Le script crée et enregistre les fichiers suivants dans votre dépôt :

learned_embeds.bin : les vecteurs d'incrustation appris correspondant à vos images d'exemple
token_identifier.txt : le jeton de substitution spécial
type_of_concept.txt : le type de concept sur lequel vous entraînez le modèle (« object » ou « style »)

Une session d'entraînement complète prend environ 1 heure sur un GPU V100 unique.

Une dernière chose avant de lancer le script. 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 <cat-toy> train"
--num_validation_images=4
--validation_steps=100

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export DATA_DIR="./cat"

accelerate launch textual_inversion.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATA_DIR \
  --learnable_property="object" \
  --placeholder_token="<cat-toy>" \
  --initializer_token="toy" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=3000 \
  --learning_rate=5.0e-04 \
  --scale_lr \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --output_dir="textual_inversion_cat" \
  --push_to_hub

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

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.load_textual_inversion("sd-concepts-library/cat-toy")
image = pipeline("A <cat-toy> train", num_inference_steps=50).images[0]
image.save("cat-train.png")

Étapes suivantes

Félicitations pour l’entraînement de votre propre modèle de Textual Inversion ! 🎉 Pour en savoir plus sur l’utilisation de votre nouveau modèle, les guides suivants peuvent vous être utiles :

Apprenez à charger des embeddings Textual Inversion et à les utiliser également en tant qu'embeddings négatifs.

Update on GitHub