Custom Diffusion

Custom Diffusion est une technique d’entraînement pour personnaliser les modèles de génération d’images. Tout comme Textual Inversion, DreamBooth et LoRA, Custom Diffusion n’exige que quelques (~4-5) images d’exemple. Cette technique fonctionne en n’entraînant que les poids des couches d’attention croisée, et elle utilise un mot spécial pour représenter le nouvel apprentissage du concept. Custom Diffusion est unique car il peut également apprendre plusieurs concepts simultanément.

Si vous entraînez sur un GPU avec une vRAM limitée, vous devriez essayer d’activer xFormers avec --enable_xformers_memory_efficient_attention pour un entraînement plus rapide avec des besoins en vRAM réduits (16 Go). Pour économiser encore plus de mémoire, ajoutez --set_grads_to_none dans l’argument d’entraînement pour définir les gradients sur None au lieu de zéro (cette option peut causer certains problèmes, donc si vous en rencontrez, essayez de retirer ce paramètre).

Ce guide examine le script train_custom_diffusion.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 :

cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval

🤗 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 contient tous les paramètres pour vous aider à personnaliser votre session d’entraînement. Ils se trouvent dans la fonction parse_args(). La fonction est fournie avec des valeurs par défaut, mais vous pouvez également définir vos propres valeurs dans la commande d’entraînement si vous le souhaitez.

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

accelerate launch train_custom_diffusion.py \
  --resolution=256

De nombreux paramètres de base sont décrits dans le guide d’entraînement DreamBooth, ce guide se concentre donc sur les paramètres spécifiques à Custom Diffusion :

--freeze_model : fige les paramètres de clé et de valeur dans la couche d'attention croisée ; la valeur par défaut est crossattn_kv, mais vous pouvez la définir sur crossattn pour entraîner tous les paramètres de la couche d'attention croisée
--concepts_list : pour apprendre plusieurs concepts, fournissez un chemin vers un fichier JSON contenant les concepts
--modifier_token : un mot spécial utilisé pour représenter le concept appris
--initializer_token : un mot spécial utilisé pour initialiser les embeddings du modifier_token

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.

De nombreux paramètres de la perte de préservation antérieure sont décrits dans le guide d'entraînement DreamBooth.

Régularisation

Custom Diffusion inclut l’entraînement des images cibles avec un petit ensemble d’images réelles afin de prévenir le surapprentissage. Comme vous pouvez l’imaginer, cela peut être facile à faire lorsque vous n’entraînez que sur quelques images ! Téléchargez 200 images réelles avec clip_retrieval. Le class_prompt doit être de la même catégorie que les images cibles. Ces images sont stockées dans class_data_dir.

python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200

Pour activer la régularisation, ajoutez les paramètres suivants :

--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
--real_prior : indique s'il faut utiliser un petit ensemble d'images réelles pour éviter le surapprentissage

accelerate launch train_custom_diffusion.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="./real_reg/samples_cat" \
  --class_prompt="cat" \
  --real_prior=True \

Script d'entraînement

Une grande partie du code du script d’entraînement Custom Diffusion est similaire à celui du script DreamBooth. Ce guide se concentre donc sur le code spécifique à Custom Diffusion.

Le script d'entraînement Custom Diffusion comporte deux classes de jeu de données :

CustomDiffusionDataset : prétraite les images, les images de classe et les prompts pour l'entraînement
PromptDataset : prépare les prompts pour générer les images de classes

Ensuite, le modifier_token est ajouté au tokenizer, converti en identifiants de jetons, et les plongements de jetons sont redimensionnés pour prendre en compte le nouveau modifier_token. Les plongements du modifier_token sont ensuite initialisés avec les plongements du initializer_token. Tous les paramètres de l'encodeur de texte sont figés, à l'exception des plongements de jetons, car c'est ce que le modèle cherche à apprendre pour l'associer aux concepts.

params_to_freeze = itertools.chain(
    text_encoder.text_model.encoder.parameters(),
    text_encoder.text_model.final_layer_norm.parameters(),
    text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)

Vous devez maintenant ajouter les Custom Diffusion weights aux couches d'attention. C'est une étape essentielle pour obtenir la forme et la taille correctes des poids d'attention, ainsi que pour définir le nombre approprié de processeurs d'attention dans chaque bloc UNet.

st = unet.state_dict()
for name, _ in unet.attn_processors.items():
    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
    if name.startswith("mid_block"):
        hidden_size = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"):
        block_id = int(name[len("up_blocks.")])
        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
    elif name.startswith("down_blocks"):
        block_id = int(name[len("down_blocks.")])
        hidden_size = unet.config.block_out_channels[block_id]
    layer_name = name.split(".processor")[0]
    weights = {
        "to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
        "to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
    }
    if train_q_out:
        weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
        weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
        weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
    if cross_attention_dim is not None:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=train_kv,
            train_q_out=train_q_out,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        ).to(unet.device)
        custom_diffusion_attn_procs[name].load_state_dict(weights)
    else:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=False,
            train_q_out=False,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        )
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)

L'optimizer est initialisé pour mettre à jour les paramètres de la couche d'attention croisée :

optimizer = optimizer_class(
    itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
    if args.modifier_token is not None
    else custom_diffusion_layers.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Dans la boucle d’entraînement, il est important de ne mettre à jour que les embeddings du concept que vous essayez d’apprendre. Cela implique de mettre à zéro les gradients de tous les autres embeddings de jetons :

if args.modifier_token is not None:
    if accelerator.num_processes > 1:
        grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
    else:
        grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
    index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
    for i in range(len(modifier_token_id[1:])):
        index_grads_to_zero = index_grads_to_zero & (
            torch.arange(len(tokenizer)) != modifier_token_id[i]
        )
    grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
        index_grads_to_zero, :
    ].fill_(0)

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

Dans ce guide, vous téléchargerez et utiliserez ces images de chats d’exemple. 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 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 chats, et OUTPUT_DIR sur l’endroit où vous souhaitez enregistrer le modèle. Vous utiliserez <new1> comme mot spécial pour lier les embeddings nouvellement appris. Le script crée et enregistre des points de contrôle du modèle ainsi qu’un fichier pytorch_custom_diffusion_weights.bin dans votre dépôt.

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 un prompt de validation avec --validation_prompt. Cela est utile pour le débogage et la sauvegarde de résultats intermédiaires.

Si vous entraînez sur des visages humains, l’équipe Custom Diffusion a constaté que les paramètres suivants fonctionnent bien :

--learning_rate=5e-6

--max_train_steps peut être compris entre 1000 et 2000

--freeze_model=crossattn

utiliser au moins 15 à 20 images pour l'entraînement

concept unique

plusieurs concepts

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

concept unique

plusieurs concepts

Étapes suivantes

Félicitations pour avoir entraîné un modèle avec Custom Diffusion ! 🎉 Pour en savoir plus :

Lisez le billet de blog Multi-Concept Customization of Text-to-Image Diffusion pour en savoir plus sur les résultats expérimentaux de l'équipe Custom Diffusion.

Update on GitHub

Diffusers

Custom Diffusion

Paramètres du script

Perte de préservation des antécédents

Régularisation

Script d'entraînement

Lancer le script

Étapes suivantes