Distillation de cohérence latente

Les Latent Consistency Models (LCMs) sont capables de générer des images de haute qualité en seulement quelques étapes, ce qui constitue un grand pas en avant, car de nombreux pipelines nécessitent au moins 25 étapes. Les LCMs sont obtenus en appliquant la méthode de distillation de cohérence latente à n’importe quel modèle Stable Diffusion. Cette méthode fonctionne en appliquant une distillation guidée en une étape à l’espace latent, et en intégrant une méthode de saut d’étapes pour sauter de manière cohérente les timesteps afin d’accélérer le processus de distillation (reportez-vous aux sections 4.1, 4.2 et 4.3 du papier pour plus de détails).

Si vous entraînez sur un GPU doté d’une vRAM limitée, essayez d’activer gradient_checkpointing, gradient_accumulation_steps et mixed_precision afin de réduire la consommation mémoire et d’accélérer l’entraînement. Vous pouvez encore diminuer la consommation mémoire en activant l’attention efficace en mémoire avec xFormers et l’optimiseur 8 bits de bitsandbytes.

Ce guide examine le script train_lcm_distill_sd_wds.py afin de vous aider à mieux le connaître et à 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/consistency_distillation
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 (essayez d’activer torch.compile pour accélérer considérablement l’entraînement) :

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_lcm_distill_sd_wds.py \
  --mixed_precision="fp16"

La plupart des paramètres sont identiques à ceux du guide d’entraînement Text-to-image, vous vous concentrerez donc sur les paramètres pertinents pour la distillation de cohérence latente dans ce guide.

• --pretrained_teacher_model : le chemin vers un modèle de diffusion latente préentraîné à utiliser comme modèle enseignant
• --pretrained_vae_model_name_or_path : chemin vers un VAE pré-entraîné ; le VAE de SDXL est connu pour souffrir d'instabilité numérique, ce paramètre vous permet donc de spécifier un VAE alternatif (comme ce VAE) de madebyollin qui fonctionne en fp16)
• --w_min et --w_max : les valeurs minimale et maximale de l'échelle de guidage pour l'échantillonnage de l'échelle de guidage
• --num_ddim_timesteps : le nombre de pas de temps pour l'échantillonnage DDIM
• --loss_type : le type de perte (L2 ou Huber) à calculer pour la distillation de cohérence latente ; la perte de Huber est généralement préférée car elle est plus robuste aux valeurs aberrantes
• --huber_c : le paramètre de la perte de Huber

Script d'entraînement

Le script d'entraînement commence par créer une classe de jeu de données - Text2ImageDataset - pour prétraiter les images et créer un jeu de données d'entraînement.

def transform(example):
    image = example["image"]
    image = TF.resize(image, resolution, interpolation=transforms.InterpolationMode.BILINEAR)

    c_top, c_left, _, _ = transforms.RandomCrop.get_params(image, output_size=(resolution, resolution))
    image = TF.crop(image, c_top, c_left, resolution, resolution)
    image = TF.to_tensor(image)
    image = TF.normalize(image, [0.5], [0.5])

    example["image"] = image
    return example

Pour améliorer les performances de lecture et d'écriture de grands jeux de données stockés dans le cloud, ce script utilise le format WebDataset pour créer un pipeline de prétraitement afin d'appliquer des transformations et de créer un jeu de données et un chargeur de données pour l'entraînement. Les images sont traitées et fournies à la boucle d'entraînement sans avoir à télécharger l'intégralité du jeu de données au préalable.

processing_pipeline = [
    wds.decode("pil", handler=wds.ignore_and_continue),
    wds.rename(image="jpg;png;jpeg;webp", text="text;txt;caption", handler=wds.warn_and_continue),
    wds.map(filter_keys({"image", "text"})),
    wds.map(transform),
    wds.to_tuple("image", "text"),
]

Dans la fonction main(), tous les composants nécessaires tels que le planificateur de bruit, les tokenizers, les encodeurs de texte et le VAE sont chargés. Le UNet enseignant est également chargé ici, puis vous pouvez créer un UNet élève à partir du UNet enseignant. Le UNet élève est mis à jour par l'optimiseur pendant l'entraînement.

teacher_unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_teacher_model, subfolder="unet", revision=args.teacher_revision
)

unet = UNet2DConditionModel(**teacher_unet.config)
unet.load_state_dict(teacher_unet.state_dict(), strict=False)
unet.train()

Vous pouvez maintenant créer l'optimizer pour mettre à jour les paramètres du UNet :

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

Créer le dataset :

dataset = Text2ImageDataset(
    train_shards_path_or_url=args.train_shards_path_or_url,
    num_train_examples=args.max_train_samples,
    per_gpu_batch_size=args.train_batch_size,
    global_batch_size=args.train_batch_size * accelerator.num_processes,
    num_workers=args.dataloader_num_workers,
    resolution=args.resolution,
    shuffle_buffer_size=1000,
    pin_memory=True,
    persistent_workers=True,
)
train_dataloader = dataset.train_dataloader

Ensuite, vous êtes prêt à configurer la boucle d’entraînement et à implémenter la méthode de distillation de cohérence latente (voir l’Algorithme 1 de l’article pour plus de détails). Cette section du script s’occupe d’ajouter du bruit aux latents, d’échantillonner et de créer un embedding d’échelle de guidage, puis de prédire l’image originale à partir du bruit.

pred_x_0 = predicted_origin(
    noise_pred,
    start_timesteps,
    noisy_model_input,
    noise_scheduler.config.prediction_type,
    alpha_schedule,
    sigma_schedule,
)

model_pred = c_skip_start * noisy_model_input + c_out_start * pred_x_0

Il récupère ensuite les prédictions du modèle enseignant et les prédictions du LCM, calcule la perte, puis la rétropropage au LCM.

if args.loss_type == "l2":
    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
elif args.loss_type == "huber":
    loss = torch.mean(
        torch.sqrt((model_pred.float() - target.float()) ** 2 + args.huber_c**2) - args.huber_c
    )

Si vous souhaitez en savoir plus sur le fonctionnement de la boucle d’entraînement, consultez le tutorial Understanding pipelines, models and schedulers 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 et à commencer la distillation !

Pour ce guide, vous utiliserez --train_shards_path_or_url pour spécifier le chemin du jeu de données Conceptual Captions 12M stocké sur le Hub ici. Définissez la variable d’environnement MODEL_DIR sur le nom du modèle enseignant et OUTPUT_DIR sur l’endroit où vous souhaitez enregistrer le modèle.

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

accelerate launch train_lcm_distill_sd_wds.py \
    --pretrained_teacher_model=$MODEL_DIR \
    --output_dir=$OUTPUT_DIR \
    --mixed_precision=fp16 \
    --resolution=512 \
    --learning_rate=1e-6 --loss_type="huber" --ema_decay=0.95 --adam_weight_decay=0.0 \
    --max_train_steps=1000 \
    --max_train_samples=4000000 \
    --dataloader_num_workers=8 \
    --train_shards_path_or_url="pipe:curl -L -s https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset/resolve/main/data/{00000..01099}.tar?download=true" \
    --validation_steps=200 \
    --checkpointing_steps=200 --checkpoints_total_limit=10 \
    --train_batch_size=12 \
    --gradient_checkpointing --enable_xformers_memory_efficient_attention \
    --gradient_accumulation_steps=1 \
    --use_8bit_adam \
    --resume_from_checkpoint=latest \
    --report_to=wandb \
    --seed=453645634 \
    --push_to_hub

Une fois l’entraînement terminé, vous pouvez utiliser votre nouveau LCM pour l’inférence.

from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
import torch

unet = UNet2DConditionModel.from_pretrained("your-username/your-model", dtype=torch.float16, variant="fp16")
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, dtype=torch.float16, variant="fp16")

pipeline.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "sushi rolls in the form of panda heads, sushi platter"

image = pipeline(prompt, num_inference_steps=4, guidance_scale=1.0).images[0]

LoRA

LoRA est une technique d’entraînement permettant de réduire considérablement le nombre de paramètres entraînables. En conséquence, l’entraînement est plus rapide et il est plus facile de stocker les poids résultants, car ils sont beaucoup plus petits (~100 Mo). Utilisez le script train_lcm_distill_lora_sd_wds.py ou train_lcm_distill_lora_sdxl.wds.py pour entraîner avec LoRA.

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

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_lcm_distill_sdxl_wds.py pour entraîner un modèle SDXL avec LoRA.

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

Étapes suivantes

Félicitations pour avoir distillé un modèle LCM ! Pour en savoir plus sur le LCM, les éléments suivants peuvent vous être utiles :

• Apprenez à utiliser LCMs pour l'inférence pour le texte-vers-image, l'image-vers-image et avec des points de contrôle LoRA.
• Lisez le billet de blog SDXL en 4 étapes avec des Latent Consistency LoRAs pour en savoir plus sur les SDXL LCM-LoRA pour une inférence ultra-rapide, les comparaisons de qualité, les benchmarks, et plus encore.