ControlNet

Les modèles ControlNet sont des adaptateurs entraînés sur la base d'un autre modèle pré-entraîné. Ils permettent un contrôle plus fin de la génération d'images en conditionnant le modèle avec une image d'entrée supplémentaire. Cette image d'entrée peut être un contour canny, une carte de profondeur, une pose humaine, et bien d'autres.

Si vous entraînez sur un GPU doté d’une vRAM limitée, vous devriez essayer d’activer les paramètres gradient_checkpointing, gradient_accumulation_steps 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 d’entraînement train_controlnet.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/controlnet
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 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_controlnet.py \
  --mixed_precision="fp16"

De nombreux paramètres de base et importants sont décrits dans le guide d’entraînement Text-to-image, ce guide se concentre donc uniquement sur les paramètres pertinents pour ControlNet :

• --max_train_samples : le nombre d’exemples d’entraînement ; vous pouvez le réduire pour accélérer l’entraînement, mais si vous souhaitez diffuser des jeux de données très volumineux, vous devrez inclure ce paramètre ainsi que le paramètre --streaming dans votre commande d’entraînement
• --gradient_accumulation_steps : nombre d'étapes de mise à jour à accumuler avant la rétropropagation ; cela permet d'entraîner avec une taille de lot supérieure à ce que la mémoire de votre GPU peut généralement gérer

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_controlnet.py \
  --snr_gamma=5.0

Script d'entraînement

Comme pour les paramètres du script, un aperçu général du script d’entraînement est fourni dans le guide d’entraînement Text-to-image. En revanche, ce guide se concentre sur les parties pertinentes du script ControlNet.

Le script d’entraînement comporte une fonction make_train_dataset pour prétraiter le jeu de données avec des transformations d’images et la tokenisation des légendes. Vous constaterez qu’en plus de la tokenisation habituelle des légendes et des transformations d’images, le script inclut également des transformations pour l’image de conditionnement.

Si vous diffusez un jeu de données sur un TPU, les performances peuvent être limitées par la bibliothèque 🤗 Datasets, qui n’est pas optimisée pour les images. Pour garantir un débit maximal, il est recommandé d’explorer d’autres formats de jeux de données tels que WebDataset, TorchData et TensorFlow Datasets.

conditioning_image_transforms = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution),
        transforms.ToTensor(),
    ]
)

Dans la fonction main(), vous trouverez le code de chargement du tokenizer, de l'encodeur de texte, du scheduler et des modèles. C'est également ici que le modèle ControlNet est chargé, soit à partir de poids existants, soit initialisé aléatoirement à partir d'un UNet :

if args.controlnet_model_name_or_path:
    logger.info("Loading existing controlnet weights")
    controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path)
else:
    logger.info("Initializing controlnet weights from unet")
    controlnet = ControlNetModel.from_unet(unet)

L'optimizer est configuré pour mettre à jour les paramètres de ControlNet :

params_to_optimize = controlnet.parameters()
optimizer = optimizer_class(
    params_to_optimize,
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Enfin, dans la boucle d'entraînement, les embeddings de texte de conditionnement et l'image sont transmis aux blocs inférieurs et intermédiaires du modèle ControlNet :

encoder_hidden_states = text_encoder(batch["input_ids"])[0]
controlnet_image = batch["conditioning_pixel_values"].to(dtype=weight_dtype)

down_block_res_samples, mid_block_res_sample = controlnet(
    noisy_latents,
    timesteps,
    encoder_hidden_states=encoder_hidden_states,
    controlnet_cond=controlnet_image,
    return_dict=False,
)

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

Ce guide utilise le jeu de données fusing/fill50k, mais n'oubliez pas 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).

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 OUTPUT_DIR sur l’emplacement où vous souhaitez enregistrer le modèle.

Téléchargez les images suivantes pour conditionner votre entraînement avec :

wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_2.png

Une dernière chose avant de lancer le script ! Selon le GPU dont vous disposez, il peut être nécessaire d’activer certaines optimisations pour entraîner un ControlNet. La configuration par défaut de ce script nécessite environ 38 Go de vRAM. Si vous entraînez sur plus d’un GPU, ajoutez le paramètre --multi_gpu à la commande accelerate launch.

pip install bitsandbytes

accelerate launch train_controlnet.py \
  --gradient_checkpointing \
  --use_8bit_adam \

accelerate launch train_controlnet.py \
  --use_8bit_adam \
  --gradient_checkpointing \
  --enable_xformers_memory_efficient_attention \
  --set_grads_to_none \

accelerate config

compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 4
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: false
  zero_stage: 2
distributed_type: DEEPSPEED

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/save/model"

accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
 --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --push_to_hub

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

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch

controlnet = ControlNetModel.from_pretrained("path/to/controlnet", dtype=torch.float16)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
    "path/to/base/model", controlnet=controlnet, dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

control_image = load_image("./conditioning_image_1.png")
prompt = "pale golden rod circle with old lace background"

generator = torch.manual_seed(0)
image = pipeline(prompt, num_inference_steps=20, generator=generator, image=control_image).images[0]
image.save("./output.png")

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_controlnet_sdxl.py pour entraîner un adaptateur ControlNet pour le modèle SDXL.

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

Étapes suivantes

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

• Apprenez à utiliser un ControlNet pour l'inférence sur une variété de tâches.