Utilisation de LoRA pour un affinage de diffusion stable et efficace

Utilisation de LoRA pour une diffusion stable et efficace.

LoRA : Low-Rank Adaptation of Large Language Models est une nouvelle technique introduite par des chercheurs de Microsoft pour résoudre le problème de l’adaptation fine des modèles de langage volumineux. Les modèles puissants avec des milliards de paramètres, tels que GPT-3, sont prohibitivement coûteux à adapter pour les tâches ou les domaines particuliers. LoRA propose de figer les poids du modèle pré-entraîné et d’injecter des couches entraînables (matrices de décomposition de rang) dans chaque bloc transformateur. Cela réduit considérablement le nombre de paramètres entraînables et les besoins en mémoire GPU, car les gradients n’ont pas besoin d’être calculés pour la plupart des poids du modèle. Les chercheurs ont constaté qu’en se concentrant sur les blocs d’attention du transformateur des modèles de langage volumineux, la qualité de l’adaptation fine avec LoRA était comparable à celle de l’adaptation fine du modèle complet, tout en étant beaucoup plus rapide et nécessitant moins de calcul.

LoRA pour les diffuseurs 🧨

Même si LoRA a été initialement proposé pour les modèles de langage volumineux et démontré sur des blocs transformateurs, la technique peut également être appliquée ailleurs. Dans le cas de l’adaptation fine par diffusion stable, LoRA peut être appliqué aux couches de co-attention qui relient les représentations d’image aux descriptions qui les décrivent. Les détails de la figure suivante (extraite de l’article sur la diffusion stable) ne sont pas importants, notez simplement que les blocs jaunes sont ceux chargés de construire la relation entre les représentations d’image et de texte.

À notre connaissance, Simo Ryu (@cloneofsimo) a été le premier à mettre en place une implémentation de LoRA adaptée à la diffusion stable. N’hésitez pas à consulter leur projet GitHub pour voir des exemples, des discussions et des idées intéressantes.

• Génération d’actifs 2D IA pour le développement de jeux #4
• L’état de la vision par ordinateur chez Hugging Face 🤗
• Une plongée dans les modèles Vision-Language

Auparavant, pour injecter des matrices entraînables LoRA aussi profondément dans le modèle que dans les couches de co-attention, les gens devaient pirater le code source des diffuseurs de manière imaginative (mais fragile). Si la diffusion stable nous a montré une chose, c’est que la communauté trouve toujours des moyens de plier et d’adapter les modèles à des fins créatives, et nous adorons ça ! La possibilité de manipuler les couches de co-attention pourrait être bénéfique pour de nombreuses autres raisons, telles que faciliter l’adoption de techniques d’optimisation telles que les transformateurs xFormers. D’autres projets créatifs tels que Prompt-to-Prompt pourraient également bénéficier d’un moyen facile d’accéder à ces couches, nous avons donc décidé de fournir une méthode générale aux utilisateurs. Nous avons testé cette demande de tirage depuis fin décembre et elle a été officiellement lancée avec notre diffusion hier.

Nous avons travaillé avec @cloneofsimo pour fournir une prise en charge de la formation LoRA dans les diffuseurs, à la fois pour les méthodes Dreambooth et d’adaptation fine complète ! Ces techniques offrent les avantages suivants :

• L’entraînement est bien plus rapide, comme nous l’avons déjà discuté.
• Les exigences en calcul sont plus faibles. Nous avons pu créer un modèle entièrement adapté en utilisant une carte graphique 2080 Ti avec 11 Go de VRAM !
• Les poids entraînés sont bien, bien plus petits. Étant donné que le modèle d’origine est figé et que nous injectons de nouvelles couches à entraîner, nous pouvons enregistrer les poids des nouvelles couches dans un seul fichier d’environ 3 Mo. C’est environ mille fois plus petit que la taille d’origine du modèle UNet !

Nous sommes particulièrement enthousiastes à propos de ce dernier point. Afin que les utilisateurs puissent partager leurs modèles adaptés ou dreamboothés exceptionnels, ils devaient partager une copie complète du modèle final. Les autres utilisateurs qui souhaitent les essayer doivent télécharger les poids adaptés dans leur interface utilisateur préférée, ce qui entraîne des coûts de stockage et de téléchargement massifs combinés. À ce jour, environ 1 000 modèles Dreambooth sont enregistrés dans la bibliothèque des concepts Dreambooth, et il y en a probablement beaucoup d’autres non enregistrés dans la bibliothèque.

Avec LoRA, il est désormais possible de publier un seul fichier de 3,29 Mo pour permettre aux autres d’utiliser votre modèle adapté.

(merci à @mishig25, la première personne à avoir utilisé le terme dreamboother en tant que verbe dans une conversation normale).

Adaptation fine LoRA

L’adaptation fine du modèle complet de diffusion stable était autrefois lente et difficile, c’est pourquoi des méthodes plus légères telles que Dreambooth ou l’inversion textuelle sont devenues si populaires. Avec LoRA, il est beaucoup plus facile de faire une adaptation fine d’un modèle sur un ensemble de données personnalisé.

Diffusers propose désormais un script d’adaptation fine LoRA qui peut s’exécuter avec seulement 11 Go de mémoire GPU, sans recourir à des astuces telles que des optimiseurs sur 8 bits. Voici comment vous pouvez l’utiliser pour effectuer une adaptation fine d’un modèle en utilisant l’ensemble de données Pokémon de Lambda Labs :

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="/sddata/finetune/lora/pokemon"
export HUB_MODEL_ID="pokemon-lora"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$DATASET_NAME \
  --dataloader_num_workers=8 \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-04 \
  --max_grad_norm=1 \
  --lr_scheduler="cosine" --lr_warmup_steps=0 \
  --output_dir=${OUTPUT_DIR} \
  --push_to_hub \
  --hub_model_id=${HUB_MODEL_ID} \
  --report_to=wandb \
  --checkpointing_steps=500 \
  --validation_prompt="Totoro" \
  --seed=1337

Une chose à noter est que le taux d’apprentissage est de 1e-4, beaucoup plus élevé que les taux d’apprentissage habituels pour le fine-tuning régulier (de l’ordre de ~1e-6, généralement). Il s’agit d’un tableau de bord W&B de l’exécution précédente, qui a duré environ 5 heures sur un GPU 2080 Ti (11 Go de RAM). Je n’ai pas essayé d’optimiser les hyperparamètres, alors n’hésitez pas à l’essayer vous-même ! Sayak a effectué une autre exécution sur un T4 (16 Go de RAM), voici son modèle final , et voici un espace de démonstration qui l’utilise .

Pour plus de détails sur le support de LoRA dans les diffusers, veuillez vous référer à notre documentation – elle sera toujours mise à jour avec la mise en œuvre.

Inférence

Comme nous l’avons discuté, l’un des principaux avantages de LoRA est que vous obtenez d’excellents résultats en entraînant des ordres de grandeur moins de poids que la taille du modèle d’origine. Nous avons conçu un processus d’inférence qui permet de charger les poids supplémentaires par-dessus les poids du modèle Stable Diffusion non modifié. Voyons comment cela fonctionne.

Tout d’abord, nous utiliserons l’API Hub pour déterminer automatiquement quel était le modèle de base qui a été utilisé pour le fine-tuning d’un modèle LoRA. À partir du modèle de Sayak , nous pouvons utiliser ce code :

from huggingface_hub import model_info

# Poids de LoRA ~3 Mo
model_path = "sayakpaul/sd-model-finetuned-lora-t4"

info = model_info(model_path)
model_base = info.cardData["base_model"]
print(model_base)   # CompVis/stable-diffusion-v1-4

Cet extrait de code affichera le modèle qu’il a utilisé pour le fine-tuning, qui est CompVis/stable-diffusion-v1-4. Dans mon cas, j’ai entraîné mon modèle à partir de la version 1.5 de Stable Diffusion, donc si vous exécutez le même code avec mon modèle LoRA, vous verrez que la sortie est runwayml/stable-diffusion-v1-5.

Les informations sur le modèle de base sont automatiquement renseignées par le script de fine-tuning que nous avons vu dans la section précédente, si vous utilisez l’option --push_to_hub. Cela est enregistré en tant qu’étiquette de métadonnées dans le fichier README du référentiel du modèle, comme vous pouvez le voir ici .

Après avoir déterminé le modèle de base que nous avons utilisé pour le fine-tuning avec LoRA, nous chargeons un pipeline Stable Diffusion normal. Nous allons le personnaliser avec le DPMSolverMultistepScheduler pour une inférence très rapide :

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

pipe = StableDiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

Et c’est là que la magie opère. Nous chargeons les poids de LoRA depuis le Hub par-dessus les poids réguliers du modèle, déplaçons le pipeline sur le dispositif cuda et exécutons l’inférence :

pipe.unet.load_attn_procs(model_path)
pipe.to("cuda")

image = pipe("Pokémon vert avec un visage menaçant", num_inference_steps=25).images[0]
image.save("green_pokemon.png")

Dreamboothing avec LoRA

Dreambooth vous permet d'”enseigner” de nouveaux concepts à un modèle de diffusion stable. LoRA est compatible avec Dreambooth et le processus est similaire au fine-tuning, avec quelques avantages :

• L’entraînement est plus rapide.
• Nous avons seulement besoin de quelques images du sujet que nous voulons entraîner (5 ou 10 suffisent généralement).
• Nous pouvons ajuster l’encodeur de texte, si nous le souhaitons, pour une fidélité supplémentaire au sujet.

Pour entraîner Dreambooth avec LoRA, vous devez utiliser ce script de diffusion . Veuillez consulter le fichier README , la documentation et notre billet de blog sur l’exploration des hyperparamètres pour plus de détails.

Pour une façon rapide, bon marché et facile d’entraîner vos modèles Dreambooth avec LoRA, veuillez vérifier cet espace par hysts . Vous devez le dupliquer et lui attribuer un GPU pour qu’il fonctionne rapidement. Ce processus vous évitera de devoir configurer votre propre environnement d’entraînement et vous pourrez entraîner vos modèles en quelques minutes !

Autres méthodes

La recherche d’un fine-tuning facile n’est pas nouvelle. En plus de Dreambooth, l’inversion textuelle est une autre méthode populaire qui tente d’enseigner de nouveaux concepts à un modèle de diffusion stable entraîné. L’une des principales raisons d’utiliser l’inversion textuelle est que les poids entraînés sont également petits et faciles à partager. Cependant, ils ne fonctionnent que pour un seul sujet (ou une petite poignée d’entre eux), tandis que LoRA peut être utilisé pour un fine-tuning polyvalent, ce qui signifie qu’il peut être adapté à de nouveaux domaines ou ensembles de données.

Le réglage pivotant est une méthode qui essaie de combiner l’inversion textuelle avec LoRA. Tout d’abord, vous enseignez au modèle un nouveau concept en utilisant des techniques d’inversion textuelle, ce qui permet d’obtenir un nouvel encodage de jeton pour le représenter. Ensuite, vous entraînez cet encodage de jeton en utilisant LoRA pour profiter des avantages des deux méthodes.

Nous n’avons pas encore exploré le réglage pivotant avec LoRA. Qui est prêt à relever le défi ? 🤗

We will continue to update IPGirl; if you have any questions or suggestions, please contact us!