1.1. Introduction à 🤗 Diffusers

Dans ce notebook, vous allez entraîner votre premier modèle de diffusion pour générer des images de mignons papillons 🦋. En cours de route, vous apprendrez les composants de base de la bibliothèque 🤗 Diffusers, qui fournira une bonne assise pour les applications plus avancées que nous couvrirons plus tard dans le cours.

Débutons par une vue d’ensemble de ce qu’on va faire dans ce notebook. Nous allons :

• Voir un puissant pipeline de modèles de diffusion personnalisé en action (avec des informations sur la façon de créer votre propre version).
• Créer votre propre mini-pipeline en :
  • Récapitulant les idées principales derrière les modèles de diffusion
  • Chargement de données à partir du Hub pour l’entraînement
  • Explorer comment ajouter du bruit à ces données à l’aide d’un planificateur
  • Créer et entraîner le modèle UNet
  • Rassembler les pièces du puzzle pour en faire un pipeline fonctionnel
• Éditer et exécuter un script pour initialiser des séries d’entraînement plus longues, qui gèrera
  • Entraînement multi-GPU via 🤗 Accelerate
  • Journalisation de l’expérience pour suivre les statistiques critiques
  • Téléchargement du modèle final sur le Hub d’Hugging Face

Installation des bibliothèques

Exécutez la cellule suivante pour installer la bibliothèque 🤗 Diffusers ainsi que quelques autres prérequis :

%pip install -qq -U diffusers datasets transformers accelerate ftfy pyarrow==9.0.0

Ensuite, rendez-vous sur https://huggingface.co/settings/tokens et créez un tokens d’accès avec autorisation d’écriture si vous n’en avez pas déjà un :

Vous pouvez vous connecter avec ce token en utilisant la ligne de commande (huggingface-cli login) ou en exécutant la cellule suivante :

from huggingface_hub import notebook_login

notebook_login()

Vous devez ensuite installer Git-LFS pour télécharger les checkpoints de votre modèle :

%%capture
!sudo apt -qq install git-lfs
!git config --global credential.helper store

Enfin, importons les bibliothèques que nous utiliserons et définissons quelques fonctions de confort que nous utiliserons plus tard dans le notebook :

import numpy as np
import torch
import torch.nn.functional as F
from matplotlib import pyplot as plt
from PIL import Image


def show_images(x):
    """Étant donné un lot d'images x, faire une grille et convertir en PIL"""
    x = x * 0.5 + 0.5  # On va de (-1, 1) et revenons (0, 1)
    grid = torchvision.utils.make_grid(x)
    grid_im = grid.detach().cpu().permute(1, 2, 0).clip(0, 1) * 255
    grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
    return grid_im


def make_grid(images, size=64):
    """Étant donné une liste d'images PIL, les empiler en une ligne pour faciliter la visualisation."""
    output_im = Image.new("RGB", (size * len(images), size))
    for i, im in enumerate(images):
        output_im.paste(im.resize((size, size)), (i * size, 0))
    return output_im


# Les utilisateurs de Mac peuvent avoir besoin de device = 'mps' (non testé)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

OK, nous sommes prêts !

Exemple générique d’inférence avec Dreambooth, un avant-goût de ce qui nous attend

Si vous avez un tant soit peu consulté les médias sociaux au cours des derniers mois, vous avez certainement entendu parler de Stable Diffusion. Il s’agit d’un puissant modèle de diffusion latent conditionné par le texte (ne vous inquiétez pas, nous allons apprendre ce que cela signifie). Mais il a un défaut : il ne sait pas à quoi vous ou moi ressemblons, à moins que nous soyons suffisamment célèbres pour que nos images soient répandues sur internet.

Dreambooth nous permet de créer notre propre variante de modèle avec une connaissance supplémentaire d’un visage, d’un objet ou d’un style spécifique. Le Corridor Crew a réalisé une excellente vidéo (en anglais) en utilisant cette technique pour raconter des histoires avec des personnages cohérents, ce qui est un excellent exemple de ce que cette technique peut faire :

from IPython.display import YouTubeVideo

YouTubeVideo("W4Mcuh38wyM")

Voici un exemple d’une sortie d’un modèle entraîné sur 5 photos du jouet Monsieur Patate.

Tout d’abord, nous chargeons le pipeline. Ceci télécharge les poids du modèle depuis le Hub. Étant donné que plusieurs gigaoctets de données sont téléchargés pour une démonstration d’une ligne, vous pouvez sauter cette cellule et simplement admirer la sortie de l’exemple !

from diffusers import StableDiffusionPipeline

# Consultez https://huggingface.co/sd-dreambooth-library pour découvrir de nombreux modèles provenant de la communauté
model_id = "sd-dreambooth-library/mr-potato-head"

# Chargement du pipeline
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(
    device
)

Une fois le chargement du pipeline terminé, nous pouvons générer des images avec :

prompt = "an abstract oil painting of sks mr potato head by picasso"
image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
image

✏️ À votre tour ! Essayez vous-même avec des prompts différents. Le token sks représente un identifiant unique pour le nouveau concept : que se passe-t-il si vous l’omettez ? Vous pouvez aussi expérimenter en changeant le nombre de pas d’échantillonnage (jusqu’où pouvez-vous descendre ?) et le paramètre guidance_scale, qui détermine jusqu’à quel point le modèle va essayer de correspondre au prompt.

Il se passe beaucoup de choses dans ce pipeline ! À la fin du cours, vous saurez comment tout cela fonctionne. Pour l’instant, voyons comment nous pouvons entraîner un modèle de diffusion à partir de zéro.

MVP (Minimum Viable Pipeline)

Exemple d’inférence sur les papillons

L’API de base de 🤗 Diffusers est divisée en trois composants principaux :

• Pipelines : classes de haut niveau conçues pour générer rapidement des échantillons à partir de modèles de diffusion populaires entraînés de manière conviviale.
• Models : architectures populaires pour entraîner de nouveaux modèles de diffusion, par exemple UNet.
• Schedulers : diverses techniques pour générer des images à partir du bruit pendant l’inférence ainsi que pour générer des images bruitées pour l’entraînement.

Les pipelines sont parfaits pour les utilisateurs finaux, mais si vous êtes ici pour ce cours, nous supposons que vous voulez savoir ce qui se passe sous le capot ! Dans le reste de ce notebook, nous allons donc construire notre propre pipeline capable de générer de petites images de papillons. Voici le résultat final en action :

from diffusers import DDPMPipeline

# Chargement du pipeline de papillons
butterfly_pipeline = DDPMPipeline.from_pretrained(
    "johnowhitaker/ddpm-butterflies-32px"
).to(device)

# Création de 8 images
images = butterfly_pipeline(batch_size=8).images

# Visualisation du résultat
make_grid(images)

Ce n’est peut-être pas aussi impressionnant que l’exemple de DreamBooth, mais nous entraînons notre modèle à partir de zéro avec ~0,0001% des données utilisées pour entraîner Stable Diffusion. En parlant d’entraînement, rappelez-vous que l’entraînement d’un modèle de diffusion ressemble à ceci :

• Chargement de quelques images à partir des données entraînées.
• Ajout de bruit, en différentes quantités.
• Introduction des versions bruitées des données d’entrée dans le modèle.
• Évaluation de la capacité du modèle à débruiter ces données d’entrée
• Utilisation de ces informations pour mettre à jour les poids du modèle, et répétition.

Nous allons explorer ces étapes une par une dans les prochaines parties jusqu’à ce que nous ayons une boucle d’entraînement complète, puis nous verrons comment échantillonner à partir du modèle entraîné et comment regrouper le tout dans un pipeline pour faciliter le partage. Commençons par les données.

Télécharger le jeu de données d’entraînement

Pour cet exemple, nous utilisons un jeu de données d’images provenant du Hub d’Hugging Face. Plus précisément, cette collection de 1000 images de papillons. Il s’agit d’un très petit jeu de données, c’est pourquoi nous avons aussi inclus des lignes en commentaires pour quelques options plus importantes. Si vous préférez utiliser votre propre collection d’images, vous pouvez également utiliser l’exemple de code commenté pour charger des images à partir d’un dossier.

import torchvision
from datasets import load_dataset
from torchvision import transforms

dataset = load_dataset("huggan/smithsonian_butterflies_subset", split="train")

# Ou charger des images à partir d'un dossier local
# dataset = load_dataset("imagefolder", data_dir="path/to/folder")

# Nous entraînerons sur des images carrées de 32 pixels, mais vous pouvez aussi essayer des tailles plus grandes
image_size = 32
# Vous pouvez réduire la taille de votre batch si vous manquez de mémoire GPU
batch_size = 64

# Définition les augmentations de données
preprocess = transforms.Compose(
    [
        transforms.Resize((image_size, image_size)),  # Redimensionner
        transforms.RandomHorizontalFlip(),  # Retournement aléatoire
        transforms.ToTensor(),  # Convertir en tenseur (0, 1)
        transforms.Normalize([0.5], [0.5]),  # Passage en (-1, 1)
    ]
)


def transform(examples):
    images = [preprocess(image.convert("RGB")) for image in examples["image"]]
    return {"images": images}


dataset.set_transform(transform)

# Créer un chargeur de données à partir du jeu de données pour servir les images transformées en batchs
train_dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=batch_size, shuffle=True
)

Nous pouvons saisir un batch d’images et en visualiser quelques-unes comme suit :

xb = next(iter(train_dataloader))["images"].to(device)[:8]
print("X shape:", xb.shape)
show_images(xb).resize((8 * 64, 64), resample=Image.NEAREST)

Nous nous en tenons à un petit jeu de données avec des images de 32 pixels pour que les temps d’entraînement restent raisonnables dans ce notebook.

Définir le planificateur

Notre plan d’entraînement consiste à prendre ces images d’entrée et à leur ajouter du bruit, puis à transmettre les images bruitées au modèle. Lors de l’inférence, nous utiliserons les prédictions du modèle pour supprimer le bruit de manière itérative. Dans 🤗 Diffusers, ces deux processus sont gérés par le scheduler (planificateur).

Le planificateur de bruit détermine la quantité de bruit ajoutée à différents moments. Voici comment nous pourrions créer un planificateur en utilisant les paramètres par défaut pour l’entraînement et l’échantillonnage “DDPM” (d’après l’article d’après l’article Denoising Diffusion Probabalistic Models) :

from diffusers import DDPMScheduler

noise_scheduler = DDPMScheduler(num_train_timesteps=1000)

Le papier DDPM décrit un processus de corruption qui ajoute une petite quantité de bruit à chaque pas de temps. Étant donné $x_{t-1}$ pour un certain pas de temps, nous pouvons obtenir la version suivante (légèrement plus bruyante) $x_t$ avec :

Nous prenons $x_{t-1}$, l’échelonnons de $\sqrt{1 - \beta_t}$ et ajoutons du bruit échelonné par $\beta_t$. Ce $\beta$ est défini pour chaque $t$ selon un certain planificateur et détermine la quantité de bruit ajoutée par pas de temps. Maintenant, nous ne voulons pas nécessairement faire cette opération 500 fois pour obtenir $x_{500}$, nous avons donc une autre formule pour obtenir $x_t$ pour n’importe quel t étant donné $x_0$ :

où :

La notation mathématique fait toujours peur ! Heureusement, le planificateur s’en charge pour nous. Nous pouvons tracer $\sqrt{\bar{\alpha}_t}$ (appelé sqrt_alpha_prod) et $\sqrt{(1 - \bar{\alpha}_t)}$ (appelé sqrt_one_minus_alpha_prod) pour voir comment l’entrée ($x$) et le bruit sont mis à l’échelle et mélangés à travers différents pas de temps :

plt.plot(noise_scheduler.alphas_cumprod.cpu() ** 0.5, label=r"${\sqrt{\bar{\alpha}_t}}$")
plt.plot((1 - noise_scheduler.alphas_cumprod.cpu()) ** 0.5, label=r"$\sqrt{(1 - \bar{\alpha}_t)}$")
plt.legend(fontsize="x-large");

✏️ À votre tour ! Vous pouvez explorer comment ce graphique change avec différents paramètres pour beta_start, beta_end et beta_schedule en remplaçant l’une des options commentées ci-dessous :

## Exemple avec beaucoup de bruit ajouté :
# noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_start=0.001, beta_end=0.004)

## Le planificateur cosinus pouvant s'avérer meilleur pour les images de petite taille :
# noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='squaredcos_cap_v2')

Quel que soit le planificateur que vous avez choisi, nous pouvons maintenant l’utiliser pour ajouter du bruit en différentes quantités en utilisant la fonction noise_scheduler.add_noise comme suit :

timesteps = torch.linspace(0, 999, 8).long().to(device)
noise = torch.randn_like(xb)
noisy_xb = noise_scheduler.add_noise(xb, noise, timesteps)
print("Noisy X shape", noisy_xb.shape)
show_images(noisy_xb).resize((8 * 64, 64), resample=Image.NEAREST)

Là encore, étudiez l’effet de l’utilisation de différents planificateurs et paramètres de bruit. Cette vidéo (en anglais) explique en détail certains des calculs ci-dessus et constitue une excellente introduction à certains de ces concepts.

Définir le modèle

Nous en arrivons maintenant à l’élément central : le modèle lui-même.

La plupart des modèles de diffusion utilisent des architectures qui sont des variantes d’un U-net et c’est ce que nous utiliserons ici.

En bref :

• l’image en entrée du modèle passe par plusieurs blocs de couches ResNet, chacun divisant la taille de l’image par 2
• puis elle passe à travers le même nombre de blocs qui la suréchantillonnent.
• il y a des skip connections qui relient les caractéristiques sur le chemin du sous-échantillonnage aux couches correspondantes dans le chemin du suréchantillonnage.

L’une des principales caractéristiques de ce modèle est qu’il prédit des images de la même taille que l’entrée, ce qui est exactement ce dont nous avons besoin ici.

🤗 Diffusers nous fournit une classe UNet2DModel pratique qui crée l’architecture désirée dans PyTorch.

Créons un U-net pour la taille d’image désirée. Notez que les down_block_types correspondent aux blocs de sous-échantillonnage (en vert sur le diagramme ci-dessus), et que les up_block_types sont les blocs de suréchantillonnage (en rouge sur le diagramme) :

from diffusers import UNet2DModel

# Création d'un modèle
model = UNet2DModel(
    sample_size=image_size,  # la résolution de l'image cible
    in_channels=3,  # le nombre de canaux d'entrée, 3 pour les images RVB
    out_channels=3,  # le nombre de canaux de sortie
    layers_per_block=2,  # le nombre de couches ResNet à utiliser par bloc UNet
    block_out_channels=(64, 128, 128, 256),  # Plus de canaux -> plus de paramètres
    down_block_types=(
        "DownBlock2D",  # un bloc de sous-échantillonnage ResNet standard
        "DownBlock2D",
        "AttnDownBlock2D",  # un bloc de sous-échantillonnage ResNet avec auto-attention spatiale
        "AttnDownBlock2D",
    ),
    up_block_types=(
        "AttnUpBlock2D",
        "AttnUpBlock2D",  # un bloc de suréchantillonnage ResNet avec auto-attention spatiale
        "UpBlock2D",
        "UpBlock2D",  # un bloc de suréchantillonnage ResNet standard
    ),
)
model.to(device)

Lorsque vous traitez des données d’entrée en haute résolution, vous pouvez utiliser davantage de blocs descendants et ascendants, et ne conserver les couches d’attention que pour les couches de résolution les plus basses (inférieures) afin de réduire l’utilisation de la mémoire. Nous verrons plus tard comment vous pouvez expérimenter pour trouver les meilleurs paramètres pour votre cas d’utilisation.

Nous pouvons vérifier que le passage d’un batch de données et de pas de temps aléatoires produit une sortie de même forme que les données d’entrée :

with torch.no_grad():
    model_prediction = model(noisy_xb, timesteps).sample
model_prediction.shape

Dans la section suivante, nous verrons comment entraîner ce modèle.

Créer une boucle d’entraînement

Il est temps d’entraîner ! Voici une boucle d’optimisation typique dans PyTorch, où nous parcourons les données batch par batch et mettons à jour les paramètres de notre modèle à chaque étape à l’aide d’un optimiseur, ici, l’optimiseur AdamW avec un taux d’apprentissage de 0,0004.

Pour chaque batch de données, nous

• échantillonnons des pas de temps aléatoires
• bruitons les données en conséquence
• transmettons les données bruitées au modèle
• comparons les prédictions du modèle avec la cible (c’est-à-dire le bruit dans ce cas) en utilisant l’erreur quadratique moyenne comme fonction de perte
• mettons à jour les paramètres du modèle via loss.backward() et optimizer.step().

Au cours de ce processus, nous enregistrons aussi les pertes au fil du temps pour un tracé ultérieur.

NB : ce code prend près de 10 minutes à exécuter. N’hésitez pas à sauter ces deux cellules et à utiliser le modèle pré-entraîné si vous êtes pressé. Vous pouvez également étudier comment la réduction du nombre de canaux dans chaque couche via la définition du modèle ci-dessus peut accélérer les choses.

L’exemple officiel d’entraînement de 🤗 Diffusers entraîne un modèle plus grand sur ce jeu de données à une résolution plus élevée, et constitue une bonne référence pour ce à quoi ressemble une boucle d’entraînement moins minimale :

# Définir le planificateur de bruit
noise_scheduler = DDPMScheduler(
    num_train_timesteps=1000, beta_schedule="squaredcos_cap_v2"
)

# Boucle d'entraînement
optimizer = torch.optim.AdamW(model.parameters(), lr=4e-4)

losses = []

for epoch in range(30):
    for step, batch in enumerate(train_dataloader):
        clean_images = batch["images"].to(device)
        # Exemple de bruit à ajouter aux images
        noise = torch.randn(clean_images.shape).to(clean_images.device)
        bs = clean_images.shape[0]

        # Échantillonner un pas de temps aléatoire pour chaque image
        timesteps = torch.randint(
            0, noise_scheduler.num_train_timesteps, (bs,), device=clean_images.device
        ).long()

        # Ajouter du bruit aux images propres en fonction de l'ampleur du bruit à chaque étape
        noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

        # Obtenir la prédiction du modèle
        noise_pred = model(noisy_images, timesteps, return_dict=False)[0]

        # Calculer la perte
        loss = F.mse_loss(noise_pred, noise)
        loss.backward(loss)
        losses.append(loss.item())

        # Mise à jour des paramètres du modèle à l'aide de l'optimiseur
        optimizer.step()
        optimizer.zero_grad()

    if (epoch + 1) % 5 == 0:
        loss_last_epoch = sum(losses[-len(train_dataloader) :]) / len(train_dataloader)
        print(f"Epoch:{epoch+1}, loss: {loss_last_epoch}")

En traçant la perte, nous constatons que le modèle s’améliore rapidement dans un premier temps, puis continue à s’améliorer à un rythme plus lent (ce qui est plus évident si nous utilisons une échelle logarithmique, comme indiqué à droite) :

fig, axs = plt.subplots(1, 2, figsize=(12, 4))
axs[0].plot(losses)
axs[1].plot(np.log(losses))
plt.show()

Au lieu d’exécuter le code d’entraînement ci-dessus, vous pouvez utiliser le modèle du pipeline comme suit :

## Décommenter pour charger le modèle que j'ai entraîné plus tôt à la place :
# model = butterfly_pipeline.unet

Générer des images

Comment obtenir des images avec ce modèle ?

• Option 1 : Création d’un pipeline :

from diffusers import DDPMPipeline

image_pipe = DDPMPipeline(unet=model, scheduler=noise_scheduler)

pipeline_output = image_pipe()
pipeline_output.images[0]

Nous pouvons enregistrer un pipeline dans un dossier local comme suit :

image_pipe.save_pretrained("my_pipeline")

Inspection du contenu du dossier :

!ls my_pipeline/

model_index.json  scheduler  unet

Les sous-dossiers scheduler et unet contiennent tout ce qui est nécessaire pour recréer ces composants. Par exemple, dans le dossier unet vous trouverez les poids du modèle (diffusion_pytorch_model.bin) ainsi qu’un fichier de configuration qui spécifie l’architecture UNet.

!ls my_pipeline/unet/

config.json  diffusion_pytorch_model.bin

Ensemble, ces fichiers contiennent tout ce qui est nécessaire pour recréer le pipeline. Vous pouvez les télécharger manuellement sur le Hub pour partager le pipeline avec d’autres personnes, ou consulter le code pour le faire via l’API dans la section suivante.

• Option 2 : écrire une boucle d’échantillonnage

Si vous inspectez la méthode forward du pipeline, vous pourrez voir ce qui se passe lorsque nous lançons image_pipe() :

# ??image_pipe.forward

Nous commençons par un bruit aléatoire et parcourons les pas de temps de l’ordonnanceur du plus bruyant au moins bruyant, en supprimant une petite quantité de bruit à chaque étape sur la base de la prédiction du modèle :

# Point de départ aléatoire (8 images aléatoires) :
sample = torch.randn(8, 3, 32, 32).to(device)

for i, t in enumerate(noise_scheduler.timesteps):

    # Obtenir le modèle de prédiction
    with torch.no_grad():
        residual = model(sample, t).sample

    # Mise à jour de l'échantillon avec le pas
    sample = noise_scheduler.step(residual, t, sample).prev_sample

show_images(sample)

La fonction noise_scheduler.step() effectue les calculs nécessaires pour mettre à jour sample de manière appropriée. Il existe un certain nombre de méthodes d’échantillonnage. Dans l’unité suivante, nous verrons comment nous pouvons échanger un échantillonneur différent pour accélérer la génération d’images avec des modèles existants, et nous parlerons plus en détail de la théorie derrière l’échantillonnage des modèles de diffusion.

Pousser votre modèle vers le Hub

Dans l’exemple ci-dessus, nous avons enregistré notre pipeline dans un dossier local. Pour pousser notre modèle vers le Hub, nous aurons besoin d’un dépôt de modèles dans lequel nous pourrons pousser nos fichiers. Nous déterminerons le nom du dépôt à partir de l’ID du modèle que nous voulons donner à notre modèle (n’hésitez pas à remplacer le nom du modèle par votre propre choix ; il doit juste contenir votre nom d’utilisateur, ce que fait la fonction get_full_repo_name()) :

from huggingface_hub import get_full_repo_name

model_name = "sd-class-butterflies-32"
hub_model_id = get_full_repo_name(model_name)
hub_model_id

Ensuite, créer un dépôt de modèle sur le 🤗 Hub et pousser notre modèle :

from huggingface_hub import HfApi, create_repo

create_repo(hub_model_id)
api = HfApi()
api.upload_folder(
    folder_path="my_pipeline/scheduler", path_in_repo="", repo_id=hub_model_id
)
api.upload_folder(folder_path="my_pipeline/unet", path_in_repo="", repo_id=hub_model_id)
api.upload_file(
    path_or_fileobj="my_pipeline/model_index.json",
    path_in_repo="model_index.json",
    repo_id=hub_model_id,
)

La dernière chose à faire est de créer une belle carte modèle afin que notre générateur de papillons puisse être facilement trouvé sur le 🤗 Hub (n’hésitez pas à développer et à modifier la description !) :

from huggingface_hub import ModelCard

content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-image-generation
- diffusion-models-class
---

# Model Card for Unit 1 of the [Diffusion Models Class 🧨](https://github.com/huggingface/diffusion-models-class)

This model is a diffusion model for unconditional image generation of cute 🦋.

## Usage


```python
from diffusers import DDPMPipeline

pipeline = DDPMPipeline.from_pretrained('{hub_model_id}')
image = pipeline().images[0]
image
```python
"""

card = ModelCard(content)
card.push_to_hub(hub_model_id)

Maintenant que le modèle est sur le Hub, vous pouvez le télécharger de n’importe où en utilisant la méthode from_pretrained() de DDPMPipeline comme suit :

from diffusers import DDPMPipeline

image_pipe = DDPMPipeline.from_pretrained(hub_model_id)
pipeline_output = image_pipe()
pipeline_output.images[0]

Bien, ça marche !

Passer à l’échelle supérieure avec 🤗 Accelerate

Ce notebook a été conçu à des fins d’apprentissage, et en tant que tel, nous avons essayé de garder le code aussi minimal et propre que possible. Pour cette raison, nous avons omis certaines choses que vous pourriez souhaiter si vous deviez entraîner un modèle plus grand sur beaucoup plus de données, comme le support multi-GPU, la trace de la progression et des images d’exemple, la sauvegarde du gradient pour supporter des tailles de batch plus importantes, le téléchargement automatique des modèles et ainsi de suite. Heureusement, la plupart de ces fonctionnalités sont disponibles dans l’exemple de script d’entraînement ici..

Vous pouvez télécharger le fichier comme suit :

!wget https://github.com/huggingface/diffusers/raw/main/examples/unconditional_image_generation/train_unconditional.py

Ouvrez le fichier et vous verrez où le modèle est défini et quels sont les paramètres disponibles. Nous exécutons le script à l’aide de la commande suivante :

# Donnons un nom à notre nouveau modèle pour le Hub
model_name = "sd-class-butterflies-64"
hub_model_id = get_full_repo_name(model_name)

!accelerate launch train_unconditional.py \
  --dataset_name="huggan/smithsonian_butterflies_subset" \
  --resolution=64 \
  --output_dir={model_name} \
  --train_batch_size=32 \
  --num_epochs=50 \
  --gradient_accumulation_steps=1 \
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision="no"

Comme précédemment, poussons le modèle vers le Hub et créons une belle carte de modèle (et n’hésitez pas à l’éditer comme vous le souhaitez !):

create_repo(hub_model_id)
api = HfApi()
api.upload_folder(
    folder_path=f"{model_name}/scheduler", path_in_repo="", repo_id=hub_model_id
)
api.upload_folder(
    folder_path=f"{model_name}/unet", path_in_repo="", repo_id=hub_model_id
)
api.upload_file(
    path_or_fileobj=f"{model_name}/model_index.json",
    path_in_repo="model_index.json",
    repo_id=hub_model_id,
)

content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-image-generation
- diffusion-models-class
---

# Model Card for Unit 1 of the [Diffusion Models Class 🧨](https://github.com/huggingface/diffusion-models-class)

This model is a diffusion model for unconditional image generation of cute 🦋.

## Usage

```python
from diffusers import DDPMPipeline

pipeline = DDPMPipeline.from_pretrained('{hub_model_id}')
image = pipeline().images[0]
image
```python
"""

card = ModelCard(content)
card.push_to_hub(hub_model_id)

Environ 45 minutes plus tard, voici le résultat :

pipeline = DDPMPipeline.from_pretrained(hub_model_id).to(device)
images = pipeline(batch_size=8).images
make_grid(images)

✏️ À votre tour ! Essayez de trouver des paramètres d’entraînement/de modèle qui donnent de bons résultats en un minimum de temps, et partagez vos résultats avec la communauté. Fouillez dans le script pour voir si vous pouvez comprendre le code, et demandez des éclaircissements sur tout ce qui vous semble confus.

Pistes pour approndonfir

Nous espérons vous avoir donné un avant-goût de ce que vous pouvez faire avec la bibliothèque 🤗 Diffusers ! Voici quelques pistes possibles pour la suite :

• Essayez d’entraîner un modèle de diffusion inconditionnel sur un nouveau jeu de données. Points bonus si vous en créez un vous-même. Vous pouvez trouver d’excellents jeux de données d’images pour cette tâche dans l’organisation HugGan sur le Hub. Assurez-vous simplement de les sous-échantillonner si vous ne voulez pas attendre très longtemps pour que le modèle s’entraîne !
• Essayez DreamBooth pour créer votre propre pipeline de Stable Diffusion personnalisé en utilisant ce Space ou ce notebook.
• Modifiez le script d’entraînement pour explorer différents hyperparamètres UNet (nombre de couches, canaux, etc.), différents schémas de bruit, etc.
• Consultez le notebook Implémentation à partir de 0 pour une approche différente des idées fondamentales que nous avons abordées dans cette unité.

1.2. Implémentation à partir de 0

Il est parfois utile de considérer la version la plus simple possible d’une chose pour mieux en comprendre le fonctionnement. C’est ce que nous allons essayer de faire dans ce notebook, en commençant par un modèle de diffusion jouet pour voir comment les différents éléments fonctionnent, puis en examinant en quoi ils diffèrent d’une mise en œuvre plus complexe.

Nous examinerons :

• Le processus de corruption (ajouter du bruit aux données)
• Ce qu’est un UNet, et comment en implémenter un extrêmement minimal à partir de zéro
• L’entraînement au modèle de diffusion
• La théorie de l’échantillonnage

Ensuite, nous comparerons nos versions avec l’implémentation DDPM des diffuseurs, en explorant :

• Les améliorations par rapport à notre mini UNet
• Le schéma de bruit du DDPM
• Les différences dans l’objectif d’entraînement
• Le conditionnement du pas de temps
• Les approches d’échantillonnage

Ce notebook est assez approfondi, et peut être sauté en toute sécurité si vous n’êtes pas enthousiaste à l’idée d’une plongée en profondeur à partir de zéro !

Il convient également de noter que la plupart du code ici est utilisé à des fins d’illustration, et nous ne recommandons pas de l’adopter directement pour votre propre travail (à moins que vous n’essayiez d’améliorer les exemples montrés ici à des fins d’apprentissage).

Configuration et importations

!pip install -q diffusers

import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from diffusers import DDPMScheduler, UNet2DModel
from matplotlib import pyplot as plt

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')

Les données

Nous allons tester les choses avec un très petit jeu de données : MNIST. Si vous souhaitez donner au modèle un défi un peu plus difficile à relever sans rien changer d’autre, torchvision.datasets.FashionMNIST devrait faire l’affaire.

dataset = torchvision.datasets.MNIST(root="mnist/", train=True, download=True, transform=torchvision.transforms.ToTensor())

train_dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

x, y = next(iter(train_dataloader))
print('Input shape:', x.shape)
print('Labels:', y)
plt.imshow(torchvision.utils.make_grid(x)[0], cmap='Greys')

Input shape: torch.Size([8, 1, 28, 28]) 
Labels: tensor([1, 9, 7, 3, 5, 2, 1, 4])

Chaque image est un dessin en niveaux de gris de 28 par 28 pixels d’un chiffre, avec des valeurs allant de 0 à 1.

Le processus de corruption

Supposons que vous n’ayez lu aucun papier sur les modèles de diffusion, mais que vous sachiez que le processus implique l’ajout de bruit. Comment feriez-vous ?

Nous souhaitons probablement disposer d’un moyen simple de contrôler le degré de corruption. Et si nous prenions un paramètre pour la quantité de bruit à ajouter, et que nous le faisions :

noise = torch.rand_like(x)

noisy_x = (1-amount)*x + amount*noise

Si amount = 0, nous récupérons l’entrée sans aucun changement. Si le montant atteint $1$, nous récupérons du bruit sans aucune trace de l’entrée $x$. En mélangeant l’entrée avec du bruit de cette façon, nous gardons la sortie dans la même plage ($0$ à $1$).

Nous pouvons mettre cela en œuvre assez facilement (il suffit de surveiller les formes pour ne pas se faire piéger par les règles de diffusion) :

def corrupt(x, amount):
  """Corrompre l'entrée `x` en la mélangeant avec du bruit selon `amount`"""
  noise = torch.rand_like(x)
  amount = amount.view(-1, 1, 1, 1) # Trier les formes pour que la transmission fonctionne
  return x*(1-amount) + noise*amount 

Et regarder les résultats visuellement pour voir que cela fonctionne comme prévu :

# Tracer les données d'entrée
fig, axs = plt.subplots(2, 1, figsize=(12, 5))
axs[0].set_title('Input data')
axs[0].imshow(torchvision.utils.make_grid(x)[0], cmap='Greys')

# Ajouter du bruit
amount = torch.linspace(0, 1, x.shape[0]) # De gauche à droite -> plus de corruption
noised_x = corrupt(x, amount)

# Tracé de la version bruitée
axs[1].set_title('Corrupted data (-- amount increases -->)')
axs[1].imshow(torchvision.utils.make_grid(noised_x)[0], cmap='Greys')

Lorsque la quantité de bruit s’approche de 1, nos données commencent à ressembler à du bruit aléatoire pur. Mais pour la plupart des noise_amounts, vous pouvez deviner le chiffre assez bien. Pensez-vous que cela soit optimal ?

Le modèle

Nous aimerions un modèle qui prenne en compte des images bruitées de 28px et qui produise une prédiction de la même forme. Un choix populaire ici est une architecture appelée UNet. Inventé à l’origine pour les tâches de segmentation en imagerie médicale, un UNet se compose d’un “chemin de compression” par lequel les données sont comprimées et d’un “chemin d’expansion” par lequel elles s’étendent à nouveau jusqu’à la dimension d’origine (similaire à un autoencodeur), mais il comporte également des connexions de saut qui permettent aux informations et aux gradients de circuler à différents niveaux.

Certains UNets comportent des blocs complexes à chaque étape, mais pour cette petite démonstration, nous construirons un exemple minimal qui prend une image à un canal et la fait passer par trois couches convolutives sur le chemin descendant (les down_layers dans le diagramme et le code) et trois sur le chemin ascendant, avec des sauts de connexion entre les couches descendantes et ascendantes. Nous utiliserons max pooling pour le downsampling et nn.Upsample pour le upsampling plutôt que de nous appuyer sur des couches apprenantes comme les UNets plus complexes. Voici l’architecture approximative montrant le nombre de canaux dans la sortie de chaque couche :

Voici à quoi cela ressemble dans le code :

class BasicUNet(nn.Module):
    """Une mise en œuvre minimale du UNet"""
    def __init__(self, in_channels=1, out_channels=1):
        super().__init__()
        self.down_layers = torch.nn.ModuleList([ 
            nn.Conv2d(in_channels, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
        ])
        self.up_layers = torch.nn.ModuleList([
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, out_channels, kernel_size=5, padding=2), 
        ])
        self.act = nn.SiLU() # La fonction d'activation
        self.downscale = nn.MaxPool2d(2)
        self.upscale = nn.Upsample(scale_factor=2)

    def forward(self, x):
        h = []
        for i, l in enumerate(self.down_layers):
            x = self.act(l(x)) # À travers la couche et la fonction d'activation
            if i < 2: # Pour toutes les couches sauf la troisième (dernière) :
              h.append(x) # Stockage de la sortie pour la skip connexion 
              x = self.downscale(x) # Réduction d'échelle pour la couche suivante
              
        for i, l in enumerate(self.up_layers):
            if i > 0:
              x = self.upscale(x) # Upscale
              x += h.pop() # Récupération d'un résultat stocké (skip connection)
            x = self.act(l(x)) # Par le biais de la couche et de la fonction d'activation
            
        return x

Nous pouvons vérifier que la forme de la sortie est la même que celle de l’entrée, comme nous nous y attendions :

net = BasicUNet()
x = torch.rand(8, 1, 28, 28)
net(x).shape

torch.Size([8, 1, 28, 28])

Ce réseau compte un peu plus de 300 000 paramètres :

sum([p.numel() for p in net.parameters()])

309057

Vous pouvez envisager de modifier le nombre de canaux dans chaque couche ou d’intervertir les architectures si vous le souhaitez.

Entraîner le réseau

Que doit faire exactement le modèle ? Là encore, il y a plusieurs façons de procéder, mais pour cette démonstration, choisissons un cadre simple : étant donné une entrée corrompue noisy_x, le modèle doit produire sa meilleure estimation de ce à quoi ressemble l’original $x$. Nous comparerons cette valeur à la valeur réelle par le biais de l’erreur quadratique moyenne. Nous comparerons cette estimation à la valeur réelle par le biais de l’erreur quadratique moyenne.

Nous pouvons maintenant entraîner le réseau.

• Obtenir un batch de données
• Corrompre les données de manière aléatoire
• Nourrir le modèle avec ces données
• Comparer les prédictions du modèle avec les images propres pour calculer notre perte
• Mettre à jour les paramètres du modèle en conséquence.

N’hésitez pas à modifier ce modèle et à voir si vous pouvez l’améliorer !

# Chargeur de données (vous pouvez modifier la taille des batchs)
batch_size = 128
train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# Combien de fois devrions-nous passer les données en revue ?
n_epochs = 3

# Créer le réseau
net = BasicUNet()
net.to(device)

# Notre fonction de perte
loss_fn = nn.MSELoss()

# L'optimiseur
opt = torch.optim.Adam(net.parameters(), lr=1e-3) 

# Conserver une trace des pertes pour les consulter ultérieurement
losses = []

# La boucle d'entraînement
for epoch in range(n_epochs):

    for x, y in train_dataloader:

        # Obtenir des données et préparer la version corrompue
        x = x.to(device) # Data on the GPU
        noise_amount = torch.rand(x.shape[0]).to(device) # Pick random noise amounts
        noisy_x = corrupt(x, noise_amount) # Create our noisy x

        # Obtenir la prédiction du modèle
        pred = net(noisy_x)

        # Calculer la perte
        loss = loss_fn(pred, x) # Dans quelle mesure la sortie est-elle proche du véritable x "propre" ?

        # Rétropropager et mettre à jour les paramètres
        opt.zero_grad()
        loss.backward()
        opt.step()

        # Stocker la perte pour plus tard
        losses.append(loss.item())

    # Afficher la moyenne des valeurs de perte pour cette époque :
    avg_loss = sum(losses[-len(train_dataloader):])/len(train_dataloader)
    print(f'Finished epoch {epoch}. Average loss for this epoch: {avg_loss:05f}')

# Visualiser la courbe des pertes
plt.plot(losses)
plt.ylim(0, 0.1)

Finished epoch 0. Average loss for this epoch: 0.026736 
Finished epoch 1. Average loss for this epoch: 0.020692 
Finished epoch 2. Average loss for this epoch: 0.018887

Nous pouvons essayer de voir à quoi ressemblent les prédictions du modèle en saisissant un batch de données, en les corrompant à différents degrés et en visualisant ensuite les prédictions du modèle :

# Récupérer des données
x, y = next(iter(train_dataloader))
x = x[:8] # Seuls les 8 premiers sont utilisés pour faciliter le graphique

# Corruption avec une échelle de montants
amount = torch.linspace(0, 1, x.shape[0]) # De gauche à droite -> plus de corruption
noised_x = corrupt(x, amount)

# Obtenir les prédictions du modèle
with torch.no_grad():
  preds = net(noised_x.to(device)).detach().cpu()

# Graphique
fig, axs = plt.subplots(3, 1, figsize=(12, 7))
axs[0].set_title('Input data')
axs[0].imshow(torchvision.utils.make_grid(x)[0].clip(0, 1), cmap='Greys')
axs[1].set_title('Corrupted data')
axs[1].imshow(torchvision.utils.make_grid(noised_x)[0].clip(0, 1), cmap='Greys')
axs[2].set_title('Network Predictions')
axs[2].imshow(torchvision.utils.make_grid(preds)[0].clip(0, 1), cmap='Greys

Vous pouvez constater que pour les montants les plus faibles, les prédictions sont plutôt bonnes ! Mais lorsque le niveau devient très élevé, le modèle a moins d’éléments pour travailler, et lorsque nous arrivons à amount=1, il produit un désordre flou proche de la moyenne du jeu de données pour essayer de couvrir ses paris sur ce à quoi la sortie pourrait ressembler…

Échantillonnage

Si nos prédictions à des niveaux de bruit élevés ne sont pas très bonnes, comment générer des images ?

Et si nous partions d’un bruit aléatoire, que nous regardions les prédictions du modèle, mais que nous ne nous rapprochions que très peu de cette prédiction (disons, 20 % du chemin). Nous disposons alors d’une image très bruyante dans laquelle il y a peut-être un soupçon de structure, que nous pouvons introduire dans le modèle pour obtenir une nouvelle prédiction. Nous espérons que cette nouvelle prédiction est légèrement meilleure que la première (puisque notre point de départ est légèrement moins bruité) et que nous pouvons donc faire un autre petit pas avec cette nouvelle et meilleure prédiction.

Nous répétons l’opération plusieurs fois et (si tout se passe bien) nous obtenons une image ! Voici ce processus illustré en seulement 5 étapes, en visualisant l’entrée du modèle (à gauche) et les images débruitées prédites (à droite) à chaque étape. Notez que même si le modèle prédit l’image débruitée dès l’étape 1, nous ne faisons qu’une partie du chemin. Au fil des étapes, les structures apparaissent et sont affinées, jusqu’à ce que nous obtenions nos résultats finaux.

n_steps = 5
x = torch.rand(8, 1, 28, 28).to(device) # Commencer au hasard
step_history = [x.detach().cpu()]
pred_output_history = []

for i in range(n_steps):
    with torch.no_grad(): # Pas besoin de suivre les gradients pendant l'inférence
        pred = net(x) # Prédire le x0 débruité
    pred_output_history.append(pred.detach().cpu()) # Stocker les résultats du modèle pour les tracer
    mix_factor = 1/(n_steps - i) # Dans quelle mesure nous nous rapprochons de la prédiction
    x = x*(1-mix_factor) + pred*mix_factor # Déplacer une partie du chemin
    step_history.append(x.detach().cpu()) # Stocker l'étape pour le graphique

fig, axs = plt.subplots(n_steps, 2, figsize=(9, 4), sharex=True)
axs[0,0].set_title('x (model input)')
axs[0,1].set_title('model prediction')
for i in range(n_steps):
    axs[i, 0].imshow(torchvision.utils.make_grid(step_history[i])[0].clip(0, 1), cmap='Greys')
    axs[i, 1].imshow(torchvision.utils.make_grid(pred_output_history[i])[0].clip(0, 1), cmap='Greys')

Nous pouvons diviser le processus en plusieurs étapes et espérer ainsi obtenir de meilleures images :

n_steps = 40
x = torch.rand(64, 1, 28, 28).to(device)
for i in range(n_steps):
  noise_amount = torch.ones((x.shape[0], )).to(device) * (1-(i/n_steps)) # Starting high going low
  with torch.no_grad():
    pred = net(x)
  mix_factor = 1/(n_steps - i)
  x = x*(1-mix_factor) + pred*mix_factor
fig, ax = plt.subplots(1, 1, figsize=(12, 12))
ax.imshow(torchvision.utils.make_grid(x.detach().cpu(), nrow=8)[0].clip(0, 1), cmap='Greys')

Ce n’est pas génial, mais il y a des chiffres reconnaissables ! Vous pouvez expérimenter en entraînant plus longtemps (disons, 10 ou 20 époques) et en modifiant la configuration du modèle, le taux d’apprentissage, l’optimiseur, etc. N’oubliez pas non plus que fashionMNIST peut être remplacé en une ligne si vous voulez essayer un jeu de données un peu plus difficile.

Comparaison avec DDPM

Dans cette section, nous allons voir comment notre implémentation diffère de l’approche utilisée dans l’autre notebook (Introduction à Diffusers), qui est basé sur l’article de DDPM.

Nous verrons que

• Le diffuseur UNet2DModel est un peu plus avancé que notre BasicUNet
• Le processus de corruption est traité différemment
• L’objectif d’entraînement est différent, puisqu’il s’agit de prédire le bruit plutôt que l’image débruitée.
• Le modèle est conditionné sur la quantité de bruit présent via un conditionnement par pas de temps, où t est transmis comme un argument supplémentaire à la méthode forward.
• Il existe un certain nombre de stratégies d’échantillonnage différentes, qui devraient fonctionner mieux que notre version simpliste ci-dessus.

Un certain nombre d’améliorations ont été suggérées depuis la publication de l’article sur le DDPM, mais nous espérons que cet exemple est instructif en ce qui concerne les différentes décisions de conception possibles. Une fois que vous aurez lu cet article, vous pourrez vous plonger dans le document intitulé Elucidating the Design Space of Diffusion-Based Generative Models qui examine tous ces composants en détail et formule de nouvelles recommandations sur la manière d’obtenir les meilleures performances.

Si tout cela est trop technique ou intimidant, ne vous inquiétez pas ! N’hésitez pas à sauter le reste de ce notebook ou à le garder pour un jour de pluie.

L’UNet

Le modèle UNet2DModel de 🤗 Diffusers comporte un certain nombre d’améliorations par rapport à notre UNet de base ci-dessus :

• GroupNorm applique une normalisation par groupe aux entrées de chaque bloc
• Couches de dropout pour un entraînement plus doux
• Plusieurs couches de ResNet par bloc (si layers_per_block n’est pas fixé à 1)
• Attention (généralement utilisé uniquement pour les blocs à faible résolution)
• Conditionnement sur le pas de temps
• Blocs de sous-échantillonnage et de suréchantillonnage avec des paramètres pouvant être appris

Créons et inspectons un modèle UNet2DModel :

model = UNet2DModel(
    sample_size=28,           # la résolution de l'image cible
    in_channels=1,            # le nombre de canaux d'entrée, 3 pour les images RVB
    out_channels=1,           # le nombre de canaux de sortie
    layers_per_block=2,       # le nombre de couches ResNet à utiliser par bloc UNet
    block_out_channels=(32, 64, 64), # Correspondant à peu près à notre exemple UNet de base
    down_block_types=( 
        "DownBlock2D",        # un bloc de sous-échantillonnage ResNet normal
        "AttnDownBlock2D",    # un bloc de sous-échantillonnage ResNet avec auto-attention spatiale
        "AttnDownBlock2D",
    ), 
    up_block_types=(
        "AttnUpBlock2D", 
        "AttnUpBlock2D",      # un bloc de suréchantillonnage ResNet avec auto-attention spatiale
        "UpBlock2D",          # un bloc de suréchantillonnage ResNet standard
      ),
)
print(model)

UNet2DModel(
  (conv_in): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (time_proj): Timesteps()
  (time_embedding): TimestepEmbedding(
    (linear_1): Linear(in_features=32, out_features=128, bias=True)
    (act): SiLU()
    (linear_2): Linear(in_features=128, out_features=128, bias=True)
  )
  (down_blocks): ModuleList(
    (0): DownBlock2D(
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (norm1): GroupNorm(32, 32, eps=1e-05, affine=True)
          (conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=32, bias=True)
          (norm2): GroupNorm(32, 32, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
        )
        (1): ResnetBlock2D(
          (norm1): GroupNorm(32, 32, eps=1e-05, affine=True)
          (conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=32, bias=True)
          (norm2): GroupNorm(32, 32, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
        )
      )
      (downsamplers): ModuleList(
        (0): Downsample2D(
          (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
        )
      )
    )
    (1): AttnDownBlock2D(
      (attentions): ModuleList(
        (0): AttentionBlock(
          (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
          (query): Linear(in_features=64, out_features=64, bias=True)
          (key): Linear(in_features=64, out_features=64, bias=True)
          (value): Linear(in_features=64, out_features=64, bias=True)
          (proj_attn): Linear(in_features=64, out_features=64, bias=True)
        )
        (1): AttentionBlock(
          (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
          (query): Linear(in_features=64, out_features=64, bias=True)
          (key): Linear(in_features=64, out_features=64, bias=True)
          (value): Linear(in_features=64, out_features=64, bias=True)
          (proj_attn): Linear(in_features=64, out_features=64, bias=True)
        )
      )
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (norm1): GroupNorm(32, 32, eps=1e-05, affine=True)
          (conv1): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
          (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
          (conv_shortcut): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1))
        )
        (1): ResnetBlock2D(
          (norm1): GroupNorm(32, 64, eps=1e-05, affine=True)
          (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
          (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
        )
      )
      (downsamplers): ModuleList(
        (0): Downsample2D(
          (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
        )
      )
    )
    (2): AttnDownBlock2D(
      (attentions): ModuleList(
        (0): AttentionBlock(
          (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
          (query): Linear(in_features=64, out_features=64, bias=True)
          (key): Linear(in_features=64, out_features=64, bias=True)
          (value): Linear(in_features=64, out_features=64, bias=True)
          (proj_attn): Linear(in_features=64, out_features=64, bias=True)
        )
        (1): AttentionBlock(
          (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
          (query): Linear(in_features=64, out_features=64, bias=True)
          (key): Linear(in_features=64, out_features=64, bias=True)
          (value): Linear(in_features=64, out_features=64, bias=True)
          (proj_attn): Linear(in_features=64, out_features=64, bias=True)
        )
      )
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (norm1): GroupNorm(32, 64, eps=1e-05, affine=True)
          (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
          (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
        )
        (1): ResnetBlock2D(
          (norm1): GroupNorm(32, 64, eps=1e-05, affine=True)
          (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
          (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
        )
      )
    )
  )
  (up_blocks): ModuleList(
    (0): AttnUpBlock2D(
      (attentions): ModuleList(
        (0): AttentionBlock(
          (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
          (query): Linear(in_features=64, out_features=64, bias=True)
          (key): Linear(in_features=64, out_features=64, bias=True)
          (value): Linear(in_features=64, out_features=64, bias=True)
          (proj_attn): Linear(in_features=64, out_features=64, bias=True)
        )
        (1): AttentionBlock(
          (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
          (query): Linear(in_features=64, out_features=64, bias=True)
          (key): Linear(in_features=64, out_features=64, bias=True)
          (value): Linear(in_features=64, out_features=64, bias=True)
          (proj_attn): Linear(in_features=64, out_features=64, bias=True)
        )
        (2): AttentionBlock(
          (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
          (query): Linear(in_features=64, out_features=64, bias=True)
          (key): Linear(in_features=64, out_features=64, bias=True)
          (value): Linear(in_features=64, out_features=64, bias=True)
          (proj_attn): Linear(in_features=64, out_features=64, bias=True)
        )
      )
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (norm1): GroupNorm(32, 128, eps=1e-05, affine=True)
          (conv1): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
          (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
          (conv_shortcut): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))
        )
        (1): ResnetBlock2D(
          (norm1): GroupNorm(32, 128, eps=1e-05, affine=True)
          (conv1): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
          (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
          (conv_shortcut): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))
        )
        (2): ResnetBlock2D(
          (norm1): GroupNorm(32, 128, eps=1e-05, affine=True)
          (conv1): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
          (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
          (conv_shortcut): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))
        )
      )
      (upsamplers): ModuleList(
        (0): Upsample2D(
          (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        )
      )
    )
    (1): AttnUpBlock2D(
      (attentions): ModuleList(
        (0): AttentionBlock(
          (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
          (query): Linear(in_features=64, out_features=64, bias=True)
          (key): Linear(in_features=64, out_features=64, bias=True)
          (value): Linear(in_features=64, out_features=64, bias=True)
          (proj_attn): Linear(in_features=64, out_features=64, bias=True)
        )
        (1): AttentionBlock(
          (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
          (query): Linear(in_features=64, out_features=64, bias=True)
          (key): Linear(in_features=64, out_features=64, bias=True)
          (value): Linear(in_features=64, out_features=64, bias=True)
          (proj_attn): Linear(in_features=64, out_features=64, bias=True)
        )
        (2): AttentionBlock(
          (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
          (query): Linear(in_features=64, out_features=64, bias=True)
          (key): Linear(in_features=64, out_features=64, bias=True)
          (value): Linear(in_features=64, out_features=64, bias=True)
          (proj_attn): Linear(in_features=64, out_features=64, bias=True)
        )
      )
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (norm1): GroupNorm(32, 128, eps=1e-05, affine=True)
          (conv1): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
          (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
          (conv_shortcut): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))
        )
        (1): ResnetBlock2D(
          (norm1): GroupNorm(32, 128, eps=1e-05, affine=True)
          (conv1): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
          (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
          (conv_shortcut): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))
        )
        (2): ResnetBlock2D(
          (norm1): GroupNorm(32, 96, eps=1e-05, affine=True)
          (conv1): Conv2d(96, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
          (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
          (conv_shortcut): Conv2d(96, 64, kernel_size=(1, 1), stride=(1, 1))
        )
      )
      (upsamplers): ModuleList(
        (0): Upsample2D(
          (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        )
      )
    )
    (2): UpBlock2D(
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (norm1): GroupNorm(32, 96, eps=1e-05, affine=True)
          (conv1): Conv2d(96, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=32, bias=True)
          (norm2): GroupNorm(32, 32, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
          (conv_shortcut): Conv2d(96, 32, kernel_size=(1, 1), stride=(1, 1))
        )
        (1): ResnetBlock2D(
          (norm1): GroupNorm(32, 64, eps=1e-05, affine=True)
          (conv1): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=32, bias=True)
          (norm2): GroupNorm(32, 32, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
          (conv_shortcut): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))
        )
        (2): ResnetBlock2D(
          (norm1): GroupNorm(32, 64, eps=1e-05, affine=True)
          (conv1): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (time_emb_proj): Linear(in_features=128, out_features=32, bias=True)
          (norm2): GroupNorm(32, 32, eps=1e-05, affine=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (nonlinearity): SiLU()
          (conv_shortcut): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))
        )
      )
    )
  )
  (mid_block): UNetMidBlock2D(
    (attentions): ModuleList(
      (0): AttentionBlock(
        (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)
        (query): Linear(in_features=64, out_features=64, bias=True)
        (key): Linear(in_features=64, out_features=64, bias=True)
        (value): Linear(in_features=64, out_features=64, bias=True)
        (proj_attn): Linear(in_features=64, out_features=64, bias=True)
      )
    )
    (resnets): ModuleList(
      (0): ResnetBlock2D(
        (norm1): GroupNorm(32, 64, eps=1e-05, affine=True)
        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
        (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
        (dropout): Dropout(p=0.0, inplace=False)
        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (nonlinearity): SiLU()
      )
      (1): ResnetBlock2D(
        (norm1): GroupNorm(32, 64, eps=1e-05, affine=True)
        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)
        (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)
        (dropout): Dropout(p=0.0, inplace=False)
        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (nonlinearity): SiLU()
      )
    )
  )
  (conv_norm_out): GroupNorm(32, 32, eps=1e-05, affine=True)
  (conv_act): SiLU()
  (conv_out): Conv2d(32, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
)

Comme vous pouvez le constater, il y a un peu plus de choses qui se passent ! Il a également beaucoup plus de paramètres que notre BasicUNet :

sum([p.numel() for p in model.parameters()]) # 1,7M contre les ~309k paramètres du BasicUNet

1707009

Nous pouvons reproduire l’entraînement présenté ci-dessus en utilisant ce modèle à la place de notre modèle original. Nous devons passer x et le pas de temps au modèle (ici, nous passons toujours t=0 pour montrer qu’il fonctionne sans ce conditionnement de pas de temps et pour faciliter le code d’échantillonnage, mais vous pouvez également essayer d’introduire (amount*1000) pour obtenir un équivalent de pas de temps à partir du montant de la corruption). Les lignes modifiées sont indiquées par #<<< si vous souhaitez inspecter le code.

# Dataloader (vous pouvez modifier la taille du batch)
batch_size = 128
train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# Combien de fois devrions-nous passer les données en revue ?
n_epochs = 3

# Créer le réseau
net = UNet2DModel(
    sample_size=28,           # la résolution de l'image cible
    in_channels=1,            # le nombre de canaux d'entrée, 3 pour les images RVB
    out_channels=1,           # le nombre de canaux de sortie
    layers_per_block=2,       # le nombre de couches ResNet à utiliser par bloc UNet
    block_out_channels=(32, 64, 64), # Correspondant à peu près à notre exemple UNet de base
    down_block_types=( 
        "DownBlock2D",        # un bloc de sous-échantillonnage ResNet normal
        "AttnDownBlock2D",    # un bloc de sous-échantillonnage ResNet avec auto-attention spatiale
        "AttnDownBlock2D",
    ), 
    up_block_types=(
        "AttnUpBlock2D", 
        "AttnUpBlock2D",      # un bloc de suréchantillonnage ResNet avec auto-attention spatiale
        "UpBlock2D",          # un bloc de suréchantillonnage ResNet standard
      ),
)

net.to(device)

# Notre protection contre la perte
loss_fn = nn.MSELoss()

# L'optimiseur
opt = torch.optim.Adam(net.parameters(), lr=1e-3) 

# Conserver une trace des pertes pour les visualiser plus tard
losses = []

# La boucle d'entraînement
for epoch in range(n_epochs):

    for x, y in train_dataloader:

        # Obtenir des données et préparer la version corrompue
        x = x.to(device) # Data on the GPU
        noise_amount = torch.rand(x.shape[0]).to(device) # Choisir des quantités de bruit aléatoires
        noisy_x = corrupt(x, noise_amount) # Créer notre bruit x

        # Obtenir la prédiction du modèle
        pred = net(noisy_x, 0).sample #<<< En utilisant toujours le pas de temps 0, en ajoutant .sample

        # Calculer la perte
        loss = loss_fn(pred, x) # Dans quelle mesure la sortie est-elle proche du véritable x "propre" ?

        # Rétropropager et mettre à jour les paramètres
        opt.zero_grad()
        loss.backward()
        opt.step()

        # Stocker la perte pour plus tard
        losses.append(loss.item())

    # Afficher la moyenne des valeurs de perte pour cette époque :
    avg_loss = sum(losses[-len(train_dataloader):])/len(train_dataloader)
    print(f'Finished epoch {epoch}. Average loss for this epoch: {avg_loss:05f}')

# Graphique
fig, axs = plt.subplots(1, 2, figsize=(12, 5))

# Perte
axs[0].plot(losses)
axs[0].set_ylim(0, 0.1)
axs[0].set_title('Loss over time')

# Échantillons
n_steps = 40
x = torch.rand(64, 1, 28, 28).to(device)
for i in range(n_steps):
  noise_amount = torch.ones((x.shape[0], )).to(device) * (1-(i/n_steps)) # De haut en bas
  with torch.no_grad():
    pred = net(x, 0).sample
  mix_factor = 1/(n_steps - i)
  x = x*(1-mix_factor) + pred*mix_factor

axs[1].imshow(torchvision.utils.make_grid(x.detach().cpu(), nrow=8)[0].clip(0, 1), cmap='Greys')
axs[1].set_title('Generated Samples')

Finished epoch 0. Average loss for this epoch: 0.018925 
Finished epoch 1. Average loss for this epoch: 0.012785 
Finished epoch 2. Average loss for this epoch: 0.011694

Ces résultats sont bien meilleurs que notre première série de résultats ! Vous pouvez envisager de modifier la configuration du Unet ou de prolonger l’entraînement afin d’obtenir des performances encore meilleures.

Le processus de corruption

Le papier DDPM décrit un processus de corruption qui ajoute une petite quantité de bruit à chaque « pas de temps ». Étant donné $x_{t-1}$ pour un certain pas de temps, nous pouvons obtenir la version suivante (légèrement plus bruitée) $x_t$ avec :

Nous prenons $x_{t-1}$, l’échelonnons de $\sqrt{1 - \beta_t}$ et ajoutons du bruit échelonné de $\beta_t$. Ce $\beta$ est défini pour chaque t en fonction d’un certain planificateur, et détermine la quantité de bruit ajoutée par pas de temps. Nous ne voulons pas nécessairement faire cette opération 500 fois pour obtenir $x_{500}$, nous avons donc une autre formule pour obtenir $x_t$ pour n’importe quel t étant donné $x_0$ :

où :

noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
plt.plot(noise_scheduler.alphas_cumprod.cpu() ** 0.5, label=r"${\sqrt{\bar{\alpha}_t}}$")
plt.plot((1 - noise_scheduler.alphas_cumprod.cpu()) ** 0.5, label=r"$\sqrt{(1 - \bar{\alpha}_t)}$")
plt.legend(fontsize="x-large")

Au départ, le $x$ bruité est principalement $x$ (sqrt_alpha_prod ~= 1), mais au fil du temps, la contribution de $x$ diminue et la composante bruit augmente. Contrairement à notre mélange linéaire de $x$ et de bruit en fonction de la quantité, celui-ci devient bruyant relativement rapidement. Nous pouvons visualiser cela sur quelques données :

# Bruit d'un batch d'images pour visualiser l'effet
fig, axs = plt.subplots(3, 1, figsize=(16, 10))
xb, yb = next(iter(train_dataloader))
xb = xb.to(device)[:8]
xb = xb * 2. - 1. # Pour aller dans (-1, 1)
print('X shape', xb.shape)

# Afficher les entrées propres
axs[0].imshow(torchvision.utils.make_grid(xb[:8])[0].detach().cpu(), cmap='Greys')
axs[0].set_title('Clean X')

# Ajouter du bruit avec le planificateur
timesteps = torch.linspace(0, 999, 8).long().to(device)
noise = torch.randn_like(xb) # << NB: randn et non rand
noisy_xb = noise_scheduler.add_noise(xb, noise, timesteps)
print('Noisy X shape', noisy_xb.shape)

# Afficher la version bruyante (avec et sans coupure)
axs[1].imshow(torchvision.utils.make_grid(noisy_xb[:8])[0].detach().cpu().clip(-1, 1),  cmap='Greys')
axs[1].set_title('Noisy X (clipped to (-1, 1)')
axs[2].imshow(torchvision.utils.make_grid(noisy_xb[:8])[0].detach().cpu(),  cmap='Greys')
axs[2].set_title('Noisy X')

X shape torch.Size([8, 1, 28, 28]) 
Noisy X shape torch.Size([8, 1, 28, 28])

Une autre dynamique est en jeu : la version DDPM ajoute un bruit tiré d’une distribution gaussienne (moyenne 0, écart-type 1 de torch.randn) plutôt que le bruit uniforme entre 0 et 1 (de torch.rand) que nous avons utilisé dans notre fonction corrompue d’origine. En général, il est judicieux de normaliser également les données d’entraînement. Dans l’autre notebook, vous verrez Normalize(0.5, 0.5) dans la liste des transformations, qui fait correspondre les données de l’image de (0, 1) à (-1, 1) et qui est “suffisante” pour nos besoins. Nous ne l’avons pas fait pour ce notebook, mais la cellule de visualisation ci-dessus l’ajoute pour une mise à l’échelle et une visualisation plus précises.

Objectif d’entraînement

Dans notre exemple, le modèle tente de prédire l’image débruitée. Dans le DDPM et dans de nombreuses autres implémentations de modèles de diffusion, le modèle prédit le bruit utilisé dans le processus de corruption (avant la mise à l’échelle, donc un bruit à variance unitaire). Dans le code, cela ressemble à quelque chose comme :

noise = torch.randn_like(xb) # << NB: randn et non rand
noisy_x = noise_scheduler.add_noise(x, noise, timesteps)
model_prediction = model(noisy_x, timesteps).sample
loss = mse_loss(model_prediction, noise) # le bruit comme cible

Vous pouvez penser que prédire le bruit (à partir duquel nous pouvons déduire à quoi ressemble l’image débruitée) est équivalent à prédire directement l’image débruitée. Alors pourquoi privilégier l’une plutôt que l’autre : est-ce simplement pour des raisons de commodité mathématique ?

Il s’avère qu’il existe une autre subtilité. Nous calculons la perte à différents moments (choisis au hasard) au cours de l’entraînement. Ces différents objectifs conduiront à une “pondération implicite” différente de ces pertes, où la prédiction du bruit donne plus de poids aux niveaux de bruit plus faibles. Vous pouvez choisir des objectifs plus complexes pour modifier cette “pondération implicite des pertes”. Vous pouvez aussi choisir un calendrier de bruit qui donnera plus d’exemples à un niveau de bruit plus élevé. Vous pouvez demander au modèle de prédire une “vitesse” v, que nous définissons comme une combinaison de l’image et du bruit dépendant du niveau de bruit (voir Progressive Distillation for Fast Sampling of Diffusion Models). Il se peut que le modèle prédise le bruit, mais qu’il réduise ensuite la perte en fonction d’un facteur dépendant de la quantité de bruit, sur la base d’un peu de théorie (voir Perception Prioritized Training of Diffusion Models) ou d’expériences visant à déterminer quels niveaux de bruit sont les plus informatifs pour le modèle (voir Elucidating the Design Space of Diffusion-Based Generative Models). En résumé : le choix de l’objectif a un effet sur les performances du modèle, et des recherches sont en cours pour déterminer la “meilleure” option.

Pour l’instant, la prédiction du bruit (epsilon ou eps) est l’approche privilégiée, mais avec le temps, nous verrons probablement d’autres objectifs pris en charge dans la bibliothèque et utilisés dans différentes situations.

Conditionnement du pas de temps

Le modèle UNet2DModel prend en compte à la fois x et le pas de temps. Ce dernier est transformé en intégration et introduit dans le modèle à plusieurs endroits.

La théorie sous-jacente est qu’en donnant au modèle des informations sur le niveau de bruit, il peut mieux accomplir sa tâche. Bien qu’il soit possible d’entraîner un modèle sans ce conditionnement du pas de temps, cela semble améliorer les performances dans certains cas et la plupart des implémentations l’incluent, du moins dans la littérature actuelle.

Échantillonnage

Étant donné un modèle qui estime le bruit présent dans une entrée bruyante (ou qui prédit la version débruitée), comment produire de nouvelles images ?

Nous pourrions introduire du bruit pur et espérer que le modèle prédise une bonne image en tant que version débruitée en une seule étape. Cependant, comme nous l’avons vu dans les expériences ci-dessus, cela ne fonctionne généralement pas bien. C’est pourquoi nous procédons à un certain nombre de petites étapes basées sur la prédiction du modèle, en éliminant de manière itérative une petite partie du bruit à la fois.

La manière exacte de procéder dépend de la méthode d’échantillonnage utilisée. Nous n’entrerons pas dans la théorie trop profondément, mais les questions clés de la conception sont les suivantes :

• Quelle est l’ampleur du pas à franchir ? En d’autres termes, quel « calendrier de bruit » devez-vous suivre ?
• Utilisez-vous uniquement la prédiction actuelle du modèle pour informer l’étape de mise à jour (comme DDPM, DDIM et beaucoup d’autres) ? Évaluez-vous le modèle plusieurs fois pour estimer les gradients d’ordre supérieur en vue d’une étape plus importante et plus précise (méthodes d’ordre supérieur et certains solveurs d’EDO discrètes) ? Ou bien conservez-vous un historique des prédictions passées pour essayer de mieux informer l’étape de mise à jour actuelle (échantillonneurs linéaires multi-étapes et ancestraux) ?
• Ajoutez-vous du bruit supplémentaire (parfois appelé « churn ») pour ajouter plus de stochasticité (caractère aléatoire) au processus d’échantillonnage, ou le gardez-vous complètement déterministe ? De nombreux échantillonneurs contrôlent ce paramètre (tel que « eta » pour les échantillonneurs DDIM) afin que l’utilisateur puisse choisir.

La recherche sur les méthodes d’échantillonnage pour les modèles de diffusion évolue rapidement et de plus en plus de méthodes permettant de trouver de bonnes solutions en moins d’étapes sont proposées. Les courageux et les curieux trouveront peut-être intéressant de parcourir le code des différentes implémentations disponibles dans la bibliothèque 🤗 Diffusers ici ou de consulter la documentation qui renvoient souvent aux articles pertinents.

Conclusions

Nous espérons que ce notebook vous a permis d’aborder les modèles de diffusion sous un angle légèrement différent.

Ce notebook a été écrit pour le cours de Hugging Face par Jonathan Whitaker, et recoupe une version incluse dans son propre cours, The Generative Landscape. Consultez-le (en anglais) si vous souhaitez voir cet exemple de base étendu avec du bruit et du conditionnement de classe. Les questions ou les bugs peuvent être communiqués via GitHub issues ou via Discord. Vous pouvez également envoyer un message via Twitter à @johnowhitaker.

2.1. Finetuning et guidage

Dans ce notebook, nous allons couvrir deux approches principales pour adapter les modèles de diffusion existants :

• Avec le finetuning, nous réentraînerons les modèles existants sur de nouvelles données afin de modifier le type de résultats qu’ils produisent.
• Avec le guidage, nous prenons un modèle existant et dirigeons le processus de génération au moment de l’inférence pour un contrôle supplémentaire.

Ce que vous apprendrez :

A la fin de ce notebook, vous saurez comment :

• Créer une boucle d’échantillonnage et générer des échantillons plus rapidement à l’aide d’un nouveau planificateur
• Finetuner un modèle de diffusion existant sur de nouvelles données, y compris :
  • Utiliser l’accumulation du gradient pour contourner certains des problèmes liés aux petits batchs.
  • Enregistrer les échantillons dans Weights and Biases pendant l’entraînement pour suivre la progression (via le script d’exemple joint).
  • Sauvegarder le pipeline résultant et le télécharger sur le Hub
• Guider le processus d’échantillonnage avec des fonctions de perte supplémentaires pour ajouter un contrôle sur les modèles existants, y compris :
  • Explorer différentes approches de guidage avec une simple perte basée sur la couleur
  • Utiliser CLIP pour guider la génération à l’aide d’un prompt de texte
  • Partager une boucle d’échantillonnage personnalisée en utilisant Gradio et 🤗 Spaces.

Configuration et importations

Pour enregistrer vos modèles finetunés sur le Hub d’Hugging Face, vous devrez vous connecter avec un token qui a un accès en écriture. Le code ci-dessous vous invite à le faire et vous renvoie à la page des tokens de votre compte. Vous aurez également besoin d’un compte Weights and Biases si vous souhaitez utiliser le script d’entraînement pour enregistrer des échantillons au fur et à mesure que le modèle s’entraîne. Là encore, le code devrait vous inviter à vous connecter là où c’est nécessaire.

A part cela, la seule chose à faire est d’installer quelques dépendances, d’importer tout ce dont nous aurons besoin et de spécifier l’appareil que nous utiliserons :

!pip install -qq diffusers datasets accelerate wandb open-clip-torch

# Code pour se connecter au Hub d'Hugging Face, nécessaire pour partager les modèles
# Assurez-vous d'utiliser un *token* avec un accès WRITE (écriture)
from huggingface_hub import notebook_login

notebook_login()

Token is valid.
Your token has been saved in your configured git credential helpers (store).
Your token has been saved to /root/.huggingface/token
Login successful

import numpy as np
import torch
import torch.nn.functional as F
import torchvision
from datasets import load_dataset
from diffusers import DDIMScheduler, DDPMPipeline
from matplotlib import pyplot as plt
from PIL import Image
from torchvision import transforms
from tqdm.auto import tqdm

device = (
    "mps"
    if torch.backends.mps.is_available()
    else "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

Chargement d’un pipeline pré-entraîné

Pour commencer ce notebook, chargeons un pipeline existant et voyons ce que nous pouvons en faire :

image_pipe = DDPMPipeline.from_pretrained("google/ddpm-celebahq-256")
image_pipe.to(device);

La génération d’images est aussi simple que l’exécution de la méthode __call__ du pipeline en l’appelant comme une fonction :

images = image_pipe().images
images[0]

Sympathique, mais LENT ! Avant d’aborder les sujets principaux du jour, jetons un coup d’œil à la boucle d’échantillonnage proprement dite et voyons comment nous pouvons utiliser un échantillonneur plus sophistiqué pour l’accélérer.

Échantillonnage plus rapide avec DDIM

À chaque étape, le modèle est nourri d’une entrée bruyante et il lui est demandé de prédire le bruit (et donc une estimation de ce à quoi l’image entièrement débruitée pourrait ressembler). Au départ, ces prédictions ne sont pas très bonnes, c’est pourquoi nous décomposons le processus en plusieurs étapes. Cependant, l’utilisation de plus de 1000 étapes s’est avérée inutile, et une multitude de recherches récentes ont exploré la manière d’obtenir de bons échantillons avec le moins d’étapes possible.

Dans la bibliothèque 🤗 Diffusers, ces méthodes d’échantillonnage sont gérées par un planificateur, qui doit effectuer chaque mise à jour via la fonction step(). Pour générer une image, on commence par un bruit aléatoire $x$. Ensuite, pour chaque pas de temps dans le planificateur de bruit, nous introduisons l’entrée bruitée $x$ dans le modèle et transmettons la prédiction résultante à la fonction step(). Celle-ci renvoie une sortie avec un attribut prev_sample. “previous” parce que nous revenons en arrière dans le temps, d’un niveau de bruit élevé à un niveau de bruit faible (à l’inverse du processus de diffusion vers l’avant).

Voyons cela en action ! Tout d’abord, nous chargeons un planificateur, ici un DDIMScheduler basé sur le papier Denoising Diffusion Implicit Models qui peut donner des échantillons décents en beaucoup moins d’étapes que l’implémentation originale du DDPM :

# Créer un nouveau planificateur et définir le nombre d'étapes d'inférence
scheduler = DDIMScheduler.from_pretrained("google/ddpm-celebahq-256")
scheduler.set_timesteps(num_inference_steps=40)

Vous pouvez constater que ce modèle effectue 40 étapes au total, chaque saut équivalant à 25 étapes du programme original de 1000 étapes :

scheduler.timesteps

tensor([975, 950, 925, 900, 875, 850, 825, 800, 775, 750, 725, 700, 675, 650,
        625, 600, 575, 550, 525, 500, 475, 450, 425, 400, 375, 350, 325, 300,
        275, 250, 225, 200, 175, 150, 125, 100,  75,  50,  25,   0])

Créons 4 images aléatoires et exécutons la boucle d’échantillonnage, en visualisant à la fois le $x$ actuel et la version débruitée prédite au fur et à mesure de l’avancement du processus :

# Le point de départ aléatoire
x = torch.randn(4, 3, 256, 256).to(device)  # Batch de 4 images à 3 canaux de 256 x 256 px

# Boucle sur les pas de temps d'échantillonnage
for i, t in tqdm(enumerate(scheduler.timesteps)):

    # Préparer l'entrée du modèle
    model_input = scheduler.scale_model_input(x, t)

    # Obtenir la prédiction
    with torch.no_grad():
        noise_pred = image_pipe.unet(model_input, t)["sample"]

    # Calculer la forme que devrait prendre l'échantillon mis à jour à l'aide du planificateur
    scheduler_output = scheduler.step(noise_pred, t, x)

    # Mise à jour de x
    x = scheduler_output.prev_sample

    # Occasionnellement, afficher à la fois x et les images débruitées prédites
    if i % 10 == 0 or i == len(scheduler.timesteps) - 1:
        fig, axs = plt.subplots(1, 2, figsize=(12, 5))

        grid = torchvision.utils.make_grid(x, nrow=4).permute(1, 2, 0)
        axs[0].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
        axs[0].set_title(f"Current x (step {i})")

        pred_x0 = (
            scheduler_output.pred_original_sample
        )  # Non disponible pour tous les planificateurs
        grid = torchvision.utils.make_grid(pred_x0, nrow=4).permute(1, 2, 0)
        axs[1].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
        axs[1].set_title(f"Predicted denoised images (step {i})")
        plt.show()

Comme vous pouvez le voir, les prédictions initiales ne sont pas très bonnes, mais au fur et à mesure que le processus se poursuit, les résultats prédits deviennent de plus en plus précis. Si vous êtes curieux de savoir ce qui se passe à l’intérieur de la fonction step(), inspectez le code (bien commenté) avec :

# ??scheduler.step

Vous pouvez également insérer ce nouveau planificateur à la place du planificateur original fourni avec le pipeline, et échantillonner de la manière suivante :

image_pipe.scheduler = scheduler
images = image_pipe(num_inference_steps=40).images
images[0]

Très bien, nous pouvons maintenant obtenir des échantillons dans un délai raisonnable ! Cela devrait accélérer les choses au fur et à mesure que nous avançons dans le reste de ce notebook :)

Finetuning

Et maintenant, le plus amusant ! Étant donné ce pipeline pré-entraîné, comment pouvons-nous réentraîner le modèle pour générer des images sur la base de nouvelles données d’entraînement ?

Il s’avère que cela est presque identique à entraîner un modèle à partir de zéro (comme nous l’avons vu dans l’unité 1), sauf que nous commençons avec le modèle existant. Voyons cela en action et abordons quelques considérations supplémentaires au fur et à mesure.

Tout d’abord, le jeu de données : vous pouvez essayer ce jeu de données de visages vintage ou ces visages animés pour quelque chose de plus proche des données d’entraînement originales de ce modèle de visages. Mais pour le plaisir, utilisons plutôt le même petit jeu de données de papillons que nous avons utilisé pour nous entraîner à partir de zéro dans l’unité 1. Exécutez le code ci-dessous pour télécharger le jeu de données papillons et créer un chargeur de données à partir duquel nous pouvons échantillonner un batch d’images :

# Pas sur Colab ? Les commentaires avec #@ permettent de modifier l'interface utilisateur comme les titres ou les entrées
# mais peuvent être ignorés si vous travaillez sur une plateforme différente.

dataset_name = "huggan/smithsonian_butterflies_subset"  # @param
dataset = load_dataset(dataset_name, split="train")
image_size = 256  # @param
batch_size = 4  # @param
preprocess = transforms.Compose(
    [
        transforms.Resize((image_size, image_size)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ]
)


def transform(examples):
    images = [preprocess(image.convert("RGB")) for image in examples["image"]]
    return {"images": images}


dataset.set_transform(transform)

train_dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=batch_size, shuffle=True
)

print("Previewing batch:")
batch = next(iter(train_dataloader))
grid = torchvision.utils.make_grid(batch["images"], nrow=4)
plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5);

Considération 1 : notre taille de batch ici (4) est assez petite, puisque nous entraînons sur une grande taille d’image (256 pixels) en utilisant un modèle assez grand et que nous manquerons de RAM du GPU si nous augmentons trop la taille du batch. Vous pouvez réduire la taille de l’image pour accélérer les choses et permettre des batchs plus importants, mais ces modèles ont été conçus et entraînés à l’origine pour une génération de 256 pixels.

Passons maintenant à la boucle d’entraînement. Nous allons mettre à jour les poids du modèle pré-entraîné en fixant la cible d’optimisation à image_pipe.unet.parameters(). Le reste est presque identique à l’exemple de boucle d’entraînement de l’unité 1. Cela prend environ 10 minutes à exécuter sur Colab, c’est donc le bon moment pour prendre un café ou un thé pendant que vous attendez :

num_epochs = 2  # @param
lr = 1e-5  # 2param
grad_accumulation_steps = 2  # @param

optimizer = torch.optim.AdamW(image_pipe.unet.parameters(), lr=lr)

losses = []

for epoch in range(num_epochs):
    for step, batch in tqdm(enumerate(train_dataloader), total=len(train_dataloader)):
        clean_images = batch["images"].to(device)
        # bruit à ajouter aux images
        noise = torch.randn(clean_images.shape).to(clean_images.device)
        bs = clean_images.shape[0]

        # un pas de temps aléatoire pour chaque image
        timesteps = torch.randint(
            0,
            image_pipe.scheduler.num_train_timesteps,
            (bs,),
            device=clean_images.device,
        ).long()

        # Ajouter du bruit aux images propres en fonction de la magnitude du bruit à chaque pas de temps
        # (il s'agit du processus de diffusion vers l'avant)
        noisy_images = image_pipe.scheduler.add_noise(clean_images, noise, timesteps)

        # Obtenir la prédiction du modèle pour le bruit
        noise_pred = image_pipe.unet(noisy_images, timesteps, return_dict=False)[0]

        # Comparez la prédiction avec le bruit réel :
        loss = F.mse_loss(
            noise_pred, noise
        )  # NB : essayer de prédire le bruit (eps) pas (noisy_ims-clean_ims) ou juste (clean_ims)

        # Stocker pour un plot ultérieur
        losses.append(loss.item())

        # Mettre à jour les paramètres du modèle avec l'optimiseur sur la base de cette perte
        loss.backward(loss)

        # Accumulation des gradients
        if (step + 1) % grad_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

    print(
        f"Epoch {epoch} average loss: {sum(losses[-len(train_dataloader):])/len(train_dataloader)}"
    )

# Tracer la courbe de perte :
plt.plot(losses)

Considération 2 : notre signal de perte est extrêmement bruyant, puisque nous ne travaillons qu’avec quatre exemples à des niveaux de bruit aléatoires pour chaque étape. Ce n’est pas idéal pour l’entraînement. Une solution consiste à utiliser un taux d’apprentissage extrêmement faible pour limiter la taille de la mise à jour à chaque étape. Ce serait encore mieux si nous pouvions trouver un moyen d’obtenir les mêmes avantages qu’en utilisant une taille de batch plus importante sans que les besoins en mémoire ne montent en flèche…

Entrez dans l’accumulation des gradients. Si nous appelons loss.backward() plusieurs fois avant d’exécuter optimizer.step() et optimizer.zero_grad(), PyTorch accumule (somme) les gradients, fusionnant effectivement le signal de plusieurs batchs pour donner une seule (meilleure) estimation qui est ensuite utilisée pour mettre à jour les paramètres. Il en résulte moins de mises à jour totales, tout comme nous le verrions si nous utilisions une taille de batch plus importante. C’est quelque chose que de nombreux frameworks gèrent pour vous (par exemple, 🤗 Accelerate rend cela facile), mais il est agréable de le voir mis en œuvre à partir de zéro car il s’agit d’une technique utile pour traiter l’entraînement sous les contraintes de mémoire du GPU ! Comme vous pouvez le voir dans le code ci-dessus (après le commentaire # Gradient accumulation), il n’y a pas vraiment besoin de beaucoup de code.

✏️ À votre tour ! Voyez si vous pouvez ajouter l’accumulation des gradients à la boucle d’entraînement de l’unité 1. Comment se comporte-t-elle ? Réfléchissez à la manière dont vous pourriez ajuster le taux d’apprentissage en fonction du nombre d’étapes d’accumulation des gradients ; devrait-il rester identique à auparavant ?

Considération 3 : Cela prend encore beaucoup de temps, et afficher une mise à jour d’une ligne à chaque époque n’est pas suffisant pour nous donner une bonne idée de ce qui se passe. Nous devrions probablement :

• Générer quelques échantillons de temps en temps pour examiner visuellement la performance qualitativement au fur et à mesure que le modèle s’entraîne.
• Enregistrer des éléments tels que la perte et les générations d’échantillons pendant l’entraînement, peut-être en utilisant quelque chose comme Weights and Biases ou Tensorboard.

Nous avons créé un script rapide (finetune_model.py) qui reprend le code d’entraînement ci-dessus et y ajoute une fonctionnalité minimale de logging. Vous pouvez voir les logs d’un entraînement ci-dessous :

%wandb johnowhitaker/dm_finetune/2upaa341 # Vous aurez besoin d'un compte W&B pour que cela fonctionne - sautez si vous ne voulez pas vous connecter.

Il est amusant de voir comment les échantillons générés changent au fur et à mesure que l’entraînement progresse. Même si la perte ne semble pas s’améliorer beaucoup, on peut voir une progression du domaine original (images de chambres à coucher) vers les nouvelles données d’entraînement (wikiart). A la fin de ce notebook se trouve un code commenté pour finetuné un modèle en utilisant ce script comme alternative à l’exécution de la cellule ci-dessus.

✏️ À votre tour ! Voyez si vous pouvez modifier l’exemple officiel de script d’entraînement que nous avons vu dans l’unité 1 pour commencer avec un modèle pré-entraîné plutôt que d’entraîner à partir de zéro. Comparez-le au script minimal dont le lien figure ci-dessus ; quelles sont les fonctionnalités supplémentaires qui manquent au script minimal ? En générant quelques images avec ce modèle, nous pouvons voir que ces visages ont déjà l’air très étranges !

x = torch.randn(8, 3, 256, 256).to(device)  # Batch de 8
for i, t in tqdm(enumerate(scheduler.timesteps)):
    model_input = scheduler.scale_model_input(x, t)
    with torch.no_grad():
        noise_pred = image_pipe.unet(model_input, t)["sample"]
    x = scheduler.step(noise_pred, t, x).prev_sample
grid = torchvision.utils.make_grid(x, nrow=4)
plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5);

Considération 4 : Le finetuning peut être tout à fait imprévisible ! Si nous entraînions plus longtemps, nous pourrions voir des papillons parfaits. Mais les étapes intermédiaires peuvent être extrêmement intéressantes en elles-mêmes, surtout si vos intérêts sont plutôt artistiques ! Entraînez sur des périodes très courtes ou très longues et faites varier le taux d’apprentissage pour voir comment cela affecte les types de résultats produits par le modèle final.

Code pour finetuner un modèle en utilisant le script d’exemple minimal que nous avons utilisé sur le modèle de démonstration WikiArt

Si vous souhaitez entraîner un modèle similaire à celui que nous avons créé sur WikiArt, vous pouvez décommenter et exécuter les cellules ci-dessous. Comme cela prend un certain temps et peut épuiser la mémoire de votre GPU, nous vous conseillons de le faire après avoir parcouru le reste de ce notebook.

## Pour télécharger le script de finetuning :
# !wget https://github.com/huggingface/diffusion-models-class/raw/main/unit2/finetune_model.py

## Pour exécuter le script, entraînant le modèle de visage sur des visages vintage
## (l'idéal est d'exécuter ce script dans un terminal) :
# !python finetune_model.py --image_size 128 --batch_size 8 --num_epochs 16\
#     --grad_accumulation_steps 2 --start_model "google/ddpm-celebahq-256"\
#     --dataset_name "Norod78/Vintage-Faces-FFHQAligned" --wandb_project 'dm-finetune'\
#     --log_samples_every 100 --save_model_every 1000 --model_save_name 'vintageface'

Sauvegarde et chargement des pipelines finetunés

Maintenant que nous avons finetuné le UNet dans notre modèle de diffusion, sauvegardons-le dans un dossier local en exécutant :

image_pipe.save_pretrained("my-finetuned-model")

Comme nous l’avons vu dans l’unité 1, cela permet de sauvegarder la configuration, le modèle et le planificateur :

!ls {"my-finetuned-model"}

Ensuite, vous pouvez suivre les mêmes étapes que celles décrites dans le notebook d’introduction à 🤗 Diffusers de l’unité 1 pour pousser le modèle vers le Hub en vue d’une utilisation ultérieure :

# Code pour télécharger un pipeline sauvegardé localement vers le Hub
from huggingface_hub import HfApi, ModelCard, create_repo, get_full_repo_name

# Mise en place du repo et téléchargement des fichiers
model_name = "ddpm-celebahq-finetuned-butterflies-2epochs"  # @param Le nom que vous souhaitez lui donner sur le Hub
local_folder_name = "my-finetuned-model"  # @param Créé par le script ou par vous via image_pipe.save_pretrained('save_name')
description = "Describe your model here"  # @param
hub_model_id = get_full_repo_name(model_name)
create_repo(hub_model_id)
api = HfApi()
api.upload_folder(
    folder_path=f"{local_folder_name}/scheduler", path_in_repo="", repo_id=hub_model_id
)
api.upload_folder(
    folder_path=f"{local_folder_name}/unet", path_in_repo="", repo_id=hub_model_id
)
api.upload_file(
    path_or_fileobj=f"{local_folder_name}/model_index.json",
    path_in_repo="model_index.json",
    repo_id=hub_model_id,
)

# Ajouter une carte modèle (facultatif mais sympa !)
content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-image-generation
- diffusion-models-class
---

# Example Fine-Tuned Model for Unit 2 of the [Diffusion Models Class 🧨](https://github.com/huggingface/diffusion-models-class)

{description}

## Usage

```python
from diffusers import DDPMPipeline

pipeline = DDPMPipeline.from_pretrained('{hub_model_id}')
image = pipeline().images[0]
image
```python
"""

card = ModelCard(content)
card.push_to_hub(hub_model_id)
     

'https://huggingface.co/lewtun/ddpm-celebahq-finetuned-butterflies-2epochs/blob/main/README.md'

Félicitations, vous avez maintenant finetuné votre premier modèle de diffusion !

Pour le reste de ce notebook, nous utiliserons un modèle que nous avons finetuné à partir d’un modèle entraîné sur LSUN bedrooms environ une fois sur le WikiArt dataset. Si vous préférez, vous pouvez sauter cette cellule et utiliser le pipeline faces/butterflies que nous avons finetuné dans la section précédente ou en charger un depuis le Hub à la place :

# Chargement du pipeline pré-entraîné
pipeline_name = "johnowhitaker/sd-class-wikiart-from-bedrooms"
image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)

# Échantillon d'images avec un planificateur DDIM sur 40 étapes
scheduler = DDIMScheduler.from_pretrained(pipeline_name)
scheduler.set_timesteps(num_inference_steps=40)

# Point de départ aléatoire (batch de 8 images)
x = torch.randn(8, 3, 256, 256).to(device)

# Boucle d'échantillonnage minimale
for i, t in tqdm(enumerate(scheduler.timesteps)):
    model_input = scheduler.scale_model_input(x, t)
    with torch.no_grad():
        noise_pred = image_pipe.unet(model_input, t)["sample"]
    x = scheduler.step(noise_pred, t, x).prev_sample

# Voir les résultats
grid = torchvision.utils.make_grid(x, nrow=4)
plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5);

Considération 5 : Il est souvent difficile de savoir si le finetuné fonctionne bien, et ce que l’on entend par “bonnes performances” peut varier selon le cas d’utilisation. Par exemple, si vous finetuné un modèle conditionné par du texte comme Stable Diffusion sur un petit jeu de données, vous voudrez probablement qu’il conserve la plus grande partie de son apprentissage original afin de pouvoir comprendre des prompts arbitraires non couverts par votre nouveau jeu de données, tout en s’adaptant pour mieux correspondre au style de vos nouvelles données d’entraînement. Cela pourrait signifier l’utilisation d’un faible taux d’apprentissage avec quelque chose comme la moyenne exponentielle du modèle, comme démontré dans cet excellent article de blog sur la création d’une version Pokemon de Stable Diffusion. Dans une autre situation, vous pouvez vouloir ré-entraîner complètement un modèle sur de nouvelles données (comme notre exemple chambre → wikiart), auquel cas un taux d’apprentissage plus élevé et un entraînement plus poussé s’avèrent judicieux. Même si le graphique de la perte ne montre pas beaucoup d’amélioration, les échantillons s’éloignent clairement des données d’origine et s’orientent vers des résultats plus “artistiques”, bien qu’ils restent pour la plupart incohérents.

Ce qui nous amène à la section suivante, où nous examinons comment nous pourrions ajouter des conseils supplémentaires à un tel modèle pour mieux contrôler les résultats.

Guidage

Que faire si l’on souhaite exercer un certain contrôle sur les échantillons générés ? Par exemple, supposons que nous voulions biaiser les images générées pour qu’elles soient d’une couleur spécifique. Comment procéder ? C’est là qu’intervient le guidage, une technique qui permet d’ajouter un contrôle supplémentaire au processus d’échantillonnage.

La première étape consiste à créer notre fonction de conditionnement : une mesure (perte) que nous souhaitons minimiser. En voici une pour l’exemple de la couleur, qui compare les pixels d’une image à une couleur cible (par défaut, une sorte de sarcelle claire) et renvoie l’erreur moyenne :

def color_loss(images, target_color=(0.1, 0.9, 0.5)):
    """Étant donné une couleur cible (R, G, B), retourner une perte correspondant à la distance moyenne entre 
	les pixels de l'image et cette couleur. Par défaut, il s'agit d'une couleur sarcelle claire : (0.1, 0.9, 0.5)"""
    target = (
        torch.tensor(target_color).to(images.device) * 2 - 1
    )  # Map target color to (-1, 1)
    target = target[
        None, :, None, None
    ]  # Obtenir la forme nécessaire pour fonctionner avec les images (b, c, h, w)
    error = torch.abs(
        images - target
    ).mean()  # Différence absolue moyenne entre les pixels de l'image et la couleur cible
    return error

Ensuite, nous allons créer une version modifiée de la boucle d’échantillonnage où, à chaque étape, nous ferons ce qui suit :

• Créer une nouvelle version de x avec requires_grad = True
• Calculer la version débruitée (x0)
• Introduire la version prédite x0 dans notre fonction de perte
• Trouver le gradient de cette fonction de perte par rapport à x
• Utiliser ce gradient de conditionnement pour modifier x avant d’utiliser le planificateur, en espérant pousser x dans une direction qui conduira à une perte plus faible selon notre fonction d’orientation.

Il existe deux variantes que vous pouvez explorer. Dans la première, nous fixons requires_grad sur x après avoir obtenu notre prédiction de bruit du UNet, ce qui est plus efficace en termes de mémoire (puisque nous n’avons pas à retracer les gradients à travers le modèle de diffusion), mais donne un gradient moins précis. Dans le second cas, nous définissons d’abord requires_grad sur x, puis nous le faisons passer par l’unet et nous calculons le x0 prédit.

# Variante 1 : méthode rapide

# L'échelle de guidance détermine l'intensité de l'effet
guidance_loss_scale = 40  # Envisagez de modifier cette valeur à 5, ou à 100

x = torch.randn(8, 3, 256, 256).to(device)

for i, t in tqdm(enumerate(scheduler.timesteps)):

    # Préparer l'entrée du modèle
    model_input = scheduler.scale_model_input(x, t)

    # Prédire le bruit résiduel
    with torch.no_grad():
        noise_pred = image_pipe.unet(model_input, t)["sample"]

    # Fixer x.requires_grad à True
    x = x.detach().requires_grad_()

    # Obtenir la valeur prédite x0
    x0 = scheduler.step(noise_pred, t, x).pred_original_sample

    # Calculer la perte
    loss = color_loss(x0) * guidance_loss_scale
    if i % 10 == 0:
        print(i, "loss:", loss.item())

    # Obtenir le gradient
    cond_grad = -torch.autograd.grad(loss, x)[0]

    # Modifier x en fonction de ce gradient
    x = x.detach() + cond_grad

    # Le planificateur
    x = scheduler.step(noise_pred, t, x).prev_sample

# Voir le résultat
grid = torchvision.utils.make_grid(x, nrow=4)
im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
Image.fromarray(np.array(im * 255).astype(np.uint8))

0 loss: 27.279136657714844
10 loss: 11.286816596984863
20 loss: 10.683112144470215
30 loss: 10.942476272583008

Cette deuxième option nécessite presque le double de RAM GPU pour fonctionner, même si nous ne générons qu’un batch de quatre images au lieu de huit. Voyez si vous pouvez repérer la différence et réfléchissez à la raison pour laquelle cette méthode est plus « précise » :

# Variante 2 : définir x.requires_grad avant de calculer les prédictions du modèle

guidance_loss_scale = 40
x = torch.randn(4, 3, 256, 256).to(device)

for i, t in tqdm(enumerate(scheduler.timesteps)):

    # Définir requires_grad avant la passe avant du modèle
    x = x.detach().requires_grad_()
    model_input = scheduler.scale_model_input(x, t)

    # prédire (avec grad cette fois)
    noise_pred = image_pipe.unet(model_input, t)["sample"]

    # Obtenir la valeur prédite x0 :
    x0 = scheduler.step(noise_pred, t, x).pred_original_sample

    # Calculer la perte
    loss = color_loss(x0) * guidance_loss_scale
    if i % 10 == 0:
        print(i, "loss:", loss.item())

    # Obtenir le gradient
    cond_grad = -torch.autograd.grad(loss, x)[0]

    # Modifier x en fonction de ce gradient
    x = x.detach() + cond_grad

    # Le planificateur
    x = scheduler.step(noise_pred, t, x).prev_sample


grid = torchvision.utils.make_grid(x, nrow=4)
im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
Image.fromarray(np.array(im * 255).astype(np.uint8))

0 loss: 30.750328063964844
10 loss: 18.550724029541016
20 loss: 17.515094757080078
30 loss: 17.55681037902832

Dans la seconde variante, les besoins en mémoire sont plus importants et l’effet est moins prononcé, de sorte que vous pouvez penser qu’elle est inférieure. Cependant, les résultats sont sans doute plus proches des types d’images sur lesquels le modèle a été entraîné, et vous pouvez toujours augmenter l’échelle de guidage pour obtenir un effet plus important. L’approche que vous utiliserez dépendra en fin de compte de ce qui fonctionne le mieux sur le plan expérimental.

✏️ À votre tour ! Choisissez votre couleur préférée et recherchez ses valeurs dans l’espace RGB. Modifiez la ligne color_loss() dans la cellule ci-dessus pour recevoir ces nouvelles valeurs RGB et examinez les résultats ; correspondent-ils à ce que vous attendez ?

Guidage avec CLIP

Guider vers une couleur nous donne un peu de contrôle, mais que se passerait-il si nous pouvions simplement taper un texte décrivant ce que nous voulons ?

CLIP est un modèle créé par OpenAI qui nous permet de comparer des images à des légendes textuelles. C’est extrêmement puissant, car cela nous permet de quantifier à quel point une image correspond à un prompt. Et comme le processus est différentiable, nous pouvons l’utiliser comme fonction de perte pour guider notre modèle de diffusion !

Nous n’entrerons pas dans les détails ici. L’approche de base est la suivante :

• Enchâsser le prompt pour obtenir un enchâssement CLIP à 512 dimensions
• Pour chaque étape du processus du modèle de diffusion :
  • Créer plusieurs variantes de l’image débruitée prédite (le fait d’avoir plusieurs variantes permet d’obtenir un signal de perte plus propre).
  • Pour chacune d’entre elles, enchâsser l’image avec CLIP et comparez cet enchâssement avec celui du prompt (à l’aide d’une mesure appelée « distance du grand cercle »).
• Calculer le gradient de cette perte par rapport à l’image bruyante actuelle x et utiliser ce gradient pour modifier x avant de le mettre à jour avec le planificateur.

Pour une explication plus approfondie de CLIP, consultez cette leçon sur le sujet ou ce rapport sur le projet OpenCLIP que nous utilisons pour charger le modèle CLIP. Exécutez la cellule suivante pour charger un modèle CLIP :

import open_clip

clip_model, _, preprocess = open_clip.create_model_and_transforms(
    "ViT-B-32", pretrained="openai"
)
clip_model.to(device)

# Transformations pour redimensionner et augmenter une image + normalisation pour correspondre aux données entraînées par CLIP
tfms = torchvision.transforms.Compose(
    [
        torchvision.transforms.RandomResizedCrop(224),  # CROP aléatoire à chaque fois
        torchvision.transforms.RandomAffine(
            5
        ),  # Une augmentation aléatoire possible : biaiser l'image
        torchvision.transforms.RandomHorizontalFlip(),  # Vous pouvez ajouter des augmentations supplémentaires si vous le souhaitez
        torchvision.transforms.Normalize(
            mean=(0.48145466, 0.4578275, 0.40821073),
            std=(0.26862954, 0.26130258, 0.27577711),
        ),
    ]
)

# Et définir une fonction de perte qui prend une image, l'enchâsse et la compare avec les caractéristiques textuelles du prompt
def clip_loss(image, text_features):
    image_features = clip_model.encode_image(
        tfms(image)
    )  # Note : applique les transformations ci-dessus
    input_normed = torch.nn.functional.normalize(image_features.unsqueeze(1), dim=2)
    embed_normed = torch.nn.functional.normalize(text_features.unsqueeze(0), dim=2)
    dists = (
        input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)
    )  # Distance du grand cercle
    return dists.mean()

Une fois la fonction de perte définie, notre boucle d’échantillonnage guidé ressemble aux exemples précédents, en remplaçant color_loss() par notre nouvelle fonction de perte basée sur CLIP :

prompt = "Red Rose (still life), red flower painting"  # @param

# Explorer en changeant ça
guidance_scale = 8  # @param
n_cuts = 4  # @param

# Plus d'étapes -> plus de temps pour que le guidage ait un effet
scheduler.set_timesteps(50)

# Nous enchâssons un prompt avec CLIP comme cible
text = open_clip.tokenize([prompt]).to(device)
with torch.no_grad(), torch.cuda.amp.autocast():
    text_features = clip_model.encode_text(text)


x = torch.randn(4, 3, 256, 256).to(
    device
)  # L'utilisation de la RAM est élevée, vous ne voulez peut-être qu'une seule image à la fois.

for i, t in tqdm(enumerate(scheduler.timesteps)):

    model_input = scheduler.scale_model_input(x, t)

    # prédire le bruit résiduel
    with torch.no_grad():
        noise_pred = image_pipe.unet(model_input, t)["sample"]

    cond_grad = 0

    for cut in range(n_cuts):

        # nécessite un grad sur x
        x = x.detach().requires_grad_()

        # Obtenir le x0 prédit
        x0 = scheduler.step(noise_pred, t, x).pred_original_sample

        # Calculer la perte
        loss = clip_loss(x0, text_features) * guidance_scale

        # Obtenir le gradient (échelle par n_cuts puisque nous voulons la moyenne)
        cond_grad -= torch.autograd.grad(loss, x)[0] / n_cuts

    if i % 25 == 0:
        print("Step:", i, ", Guidance loss:", loss.item())

    # Modifier x en fonction de ce gradient
    alpha_bar = scheduler.alphas_cumprod[i]
    x = (
        x.detach() + cond_grad * alpha_bar.sqrt()
    )  # Note the additional scaling factor here!

    # Le planificateur
    x = scheduler.step(noise_pred, t, x).prev_sample


grid = torchvision.utils.make_grid(x.detach(), nrow=4)
im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
Image.fromarray(np.array(im * 255).astype(np.uint8))

Step: 0 , Guidance loss: 7.437869548797607
Step: 25 , Guidance loss: 7.174620628356934

Cela ressemble un peu à des roses ! Ce n’est pas parfait, mais si vous jouez avec les paramètres, vous pouvez obtenir des images agréables.

Si vous examinez le code ci-dessus, vous verrez que nous mettons à l’échelle le gradient de conditionnement par un facteur de alpha_bar.sqrt(). Il existe des théories sur la “bonne” manière d’échelonner ces gradients, mais en pratique, vous pouvez expérimenter. Pour certains types de guidage, vous voudrez peut-être que la plupart des effets soient concentrés dans les premières étapes, pour d’autres (par exemple, une perte de style axée sur les textures), vous préférerez peut-être qu’ils n’interviennent que vers la fin du processus de génération. Quelques programmes possibles sont présentés ci-dessous :

plt.plot([1 for a in scheduler.alphas_cumprod], label="no scaling")
plt.plot([a for a in scheduler.alphas_cumprod], label="alpha_bar")
plt.plot([a.sqrt() for a in scheduler.alphas_cumprod], label="alpha_bar.sqrt()")
plt.plot(
    [(1 - a).sqrt() for a in scheduler.alphas_cumprod], label="(1-alpha_bar).sqrt()"
)
plt.legend()
plt.title("Possible guidance scaling schedules")

Expérimentez avec différents planificateurs, échelles de guidage et toute autre astuce à laquelle vous pouvez penser (l’écrêtage des gradients dans une certaine plage est une modification populaire) pour voir jusqu’à quel point vous pouvez obtenir ce résultat ! N’oubliez pas non plus d’essayer d’intervertir d’autres modèles. Peut-être le modèle de visages que nous avons chargé au début ; pouvez-vous le guider de manière fiable pour produire un visage masculin ? Que se passe-t-il si vous combinez le guidage CLIP avec la perte de couleur que nous avons utilisée plus tôt ? Etc.

Si vous consultez quelques codes pour la diffusion guidée par CLIP en pratique, vous verrez une approche plus complexe avec une meilleure classe pour choisir des découpes aléatoires dans les images et de nombreux ajustements supplémentaires de la fonction de perte pour de meilleures performances. Avant l’apparition des modèles de diffusion conditionnés par le texte, il s’agissait du meilleur système de conversion texte-image qui soit ! La petite version de notre jouet peut encore être améliorée, mais elle capture l’idée principale : grâce au guidage et aux capacités étonnantes de CLIP, nous pouvons ajouter le contrôle du texte à un modèle de diffusion inconditionnel 🎨.

Partager une boucle d’échantillonnage personnalisée en tant que démo Gradio

Vous avez peut-être trouvé une perte amusante pour guider la génération et vous souhaitez maintenant partager avec le monde entier votre modèle finetuné et cette stratégie d’échantillonnage personnalisée…

Entrez dans Gradio. Gradio est un outil gratuit et open-source qui permet aux utilisateurs de créer et de partager facilement des modèles interactifs d’apprentissage automatique via une simple interface web. Avec Gradio, les utilisateurs peuvent construire des interfaces personnalisées pour leurs modèles d’apprentissage automatique, qui peuvent ensuite être partagés avec d’autres par le biais d’une URL unique. Il est également intégré à 🤗 Spaces, ce qui permet d’héberger facilement des démos et de les partager avec d’autres.

Nous placerons notre logique de base dans une fonction qui prend certaines entrées et produit une image en sortie. Cette fonction peut ensuite être enveloppée dans une interface simple qui permet à l’utilisateur de spécifier certains paramètres (qui sont transmis en tant qu’entrées à la fonction principale de génération). De nombreux composants sont disponibles ; pour cet exemple, nous utiliserons un curseur pour l’échelle d’orientation et un sélecteur de couleurs pour définir la couleur cible.

!pip install -q gradio

import gradio as gr
from PIL import Image, ImageColor

# La fonction qui fait le gros du travail
def generate(color, guidance_loss_scale):
    target_color = ImageColor.getcolor(color, "RGB")  # Couleur cible en RGB
    target_color = [a / 255 for a in target_color]  # Rééchelonner de (0, 255) à (0, 1)
    x = torch.randn(1, 3, 256, 256).to(device)
    for i, t in tqdm(enumerate(scheduler.timesteps)):
        model_input = scheduler.scale_model_input(x, t)
        with torch.no_grad():
            noise_pred = image_pipe.unet(model_input, t)["sample"]
        x = x.detach().requires_grad_()
        x0 = scheduler.step(noise_pred, t, x).pred_original_sample
        loss = color_loss(x0, target_color) * guidance_loss_scale
        cond_grad = -torch.autograd.grad(loss, x)[0]
        x = x.detach() + cond_grad
        x = scheduler.step(noise_pred, t, x).prev_sample
    grid = torchvision.utils.make_grid(x, nrow=4)
    im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
    im = Image.fromarray(np.array(im * 255).astype(np.uint8))
    im.save("test.jpeg")
    return im


# Voir la documentation de gradio pour les types d'entrées et de sorties disponibles.
inputs = [
    gr.ColorPicker(label="color", value="55FFAA"),  # Ajoutez ici toutes les entrées dont vous avez besoin
    gr.Slider(label="guidance_scale", minimum=0, maximum=30, value=3),
]
outputs = gr.Image(label="result")

# Et l'interface minimale
demo = gr.Interface(
    fn=generate,
    inputs=inputs,
    outputs=outputs,
    examples=[
        ["#BB2266", 3],
        ["#44CCAA", 5],  # Vous pouvez fournir des exemples d'entrées pour aider les gens à démarrer
    ],
)
demo.launch(debug=True)  # debug=True vous permet de voir les erreurs et les sorties dans Colab

Il est possible de construire des interfaces beaucoup plus compliquées, avec un style fantaisiste et un large éventail d’entrées possibles, mais pour cette démo, nous la gardons aussi simple que possible.

Les démos sur 🤗 Spaces s’exécutent par défaut sur CPU, il est donc préférable de prototyper votre interface dans Colab (comme ci-dessus) avant de la migrer. Lorsque vous êtes prêt à partager votre démo, vous devez créer un Space, mettre en place un fichier requirements.txt listant les bibliothèques que votre code utilisera, puis placer tout le code dans un fichier app.py qui définit les fonctions pertinentes et l’interface.

Heureusement pour vous, il est également possible de “dupliquer” un Space. Vous pouvez visiter le Space ici et cliquer sur “Dupliquer cet espace” pour obtenir un modèle que vous pouvez ensuite modifier pour utiliser votre propre modèle et votre propre fonction d’orientation.

Dans les paramètres, vous pouvez configurer votre Space pour qu’il fonctionne avec du matériel plus sophistiqué (qui est facturé à l’heure). Vous avez créé quelque chose d’extraordinaire et vous voulez le partager sur un meilleur matériel, mais vous n’avez pas l’argent nécessaire ? Faites-le nous savoir via Discord et nous verrons si nous pouvons vous aider !

Résumé et prochaines étapes

Nous avons couvert beaucoup de choses dans ce notebook ! Récapitulons les idées principales :

• Il est relativement facile de charger des modèles existants et de les échantillonner avec différents planificateurs
• Le finetuning ressemble à l’entraînement à partir de zéro, sauf qu’en partant d’un modèle existant, nous espérons obtenir de meilleurs résultats plus rapidement.
• Pour finetuner de grands modèles sur de grandes images, nous pouvons utiliser des astuces comme l’accumulation de gradient pour contourner les limitations de la taille des batchs.
• L’enregistrement d’échantillons d’images est important pour le finetuning, où une courbe de perte peut ne pas fournir beaucoup d’informations utiles.
• Le guidage nous permet de prendre un modèle inconditionnel et d’orienter le processus de génération sur la base d’une fonction de guidage/perte, où à chaque étape nous trouvons le gradient de la perte par rapport à l’image bruitée $x$ et l’actualisons en fonction de ce gradient avant de passer à l’étape temporelle suivante.
• Le guidage avec CLIP nous permet de contrôler des modèles inconditionnels avec du texte !

Pour mettre cela en pratique, voici quelques étapes spécifiques que vous pouvez suivre :

• Finetuné votre propre modèle et le pousser vers le Hub. Cela implique de choisir un point de départ (par exemple, un modèle entraîné sur faces, bedrooms, cats ou wikiart et un jeu de données (peut-être ces faces d’animaux ou vos propres images), puis d’entraîner soit le code de ce notebook, soit le script d’exemple (utilisation de démonstration ci-dessous).
• Explorer le guidage en utilisant votre modèle finetuné, soit en utilisant l’une des fonctions de guidage de l’exemple (color_loss ou CLIP), soit en inventant la vôtre.
• Partagez une démo basée sur ceci en utilisant Gradio, soit en modifiant le Space d’exemple pour utiliser votre propre modèle, soit en créant votre propre version personnalisée avec plus de fonctionnalités.

Nous sommes impatients de voir vos résultats sur Discord, Twitter et ailleurs 🤗 !

2.2. Modèle de diffusion conditionné par la classe

Dans ce notebook, nous allons illustrer une façon d’ajouter des informations de conditionnement à un modèle de diffusion. Plus précisément, nous allons entraîner un modèle de diffusion conditionné par la classe sur MNIST à la suite de l’exemple d’entraînement à partir de 0 de l’unité 1, où nous pouvons spécifier quel chiffre nous voulons que le modèle génère au moment de l’inférence.

Comme indiqué dans l’introduction de cette unité, il s’agit d’une des nombreuses façons d’ajouter des informations de conditionnement supplémentaires à un modèle de diffusion, et elle a été choisie pour sa relative simplicité. Tout comme le notebook de l’unité 1, ce notebook n’a qu’un but illustratif et vous pouvez l’ignorer si vous le souhaitez.

Configuration et préparation des données

!pip install -q diffusers

import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from diffusers import DDPMScheduler, UNet2DModel
from matplotlib import pyplot as plt
from tqdm.auto import tqdm

device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')

# Charger le jeu de données
dataset = torchvision.datasets.MNIST(root="mnist/", train=True, download=True, transform=torchvision.transforms.ToTensor())

# Introduire les données dans un chargeur de données (batch de taille 8 ici pour la démonstration)
train_dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

# Visualiser quelques exemples
x, y = next(iter(train_dataloader))
print('Input shape:', x.shape)
print('Labels:', y)
plt.imshow(torchvision.utils.make_grid(x)[0], cmap='Greys')

Input shape: torch.Size([8, 1, 28, 28]) 
Labels: tensor([8, 1, 5, 9, 7, 6, 2, 2])

Création d’une UNet conditionnée par la classe

La façon dont nous introduirons le conditionnement de la classe est la suivante :

• Créer un UNet2DModel standard avec quelques canaux d’entrée supplémentaires.
• Associer l’étiquette de la classe à un vecteur appris de forme (class_emb_size) via une couche d’enchâssement.
• Concaténer ces informations en tant que canaux supplémentaires pour l’entrée interne du UNet avec net_input = torch.cat((x, class_cond), 1)
• Introduire ce net_input (qui a (class_emb_size+1) canaux au total) dans l’UNet pour obtenir la prédiction finale.

Dans cet exemple, nous avons fixé la taille de class_emb_size à 4, mais c’est complètement arbitraire et vous pourriez envisager de la fixer à 1 (pour voir si cela fonctionne toujours), à 10 (pour correspondre au nombre de classes), ou de remplacer le nn.Embedding appris par un simple encodage à un coup de l’étiquette de la classe directement.

Voici à quoi ressemble l’implémentation :

class ClassConditionedUNet(nn.Module):
  def __init__(self, num_classes=10, class_emb_size=4):
    super().__init__()
    
    # La couche d'intégration associe l'étiquette de la classe à un vecteur de taille `class_emb_size`
    self.class_emb = nn.Embedding(num_classes, class_emb_size)

    # Self.model est un UNet inconditionnel avec des canaux d'entrée supplémentaires pour accepter les informations de conditionnement (l'enchâssement de la classe).
    self.model = UNet2DModel(
        sample_size=28,           # la résolution de l'image cible
        in_channels=1 + class_emb_size, # Canaux d'entrée supplémentaires pour la classe conditionnée
        out_channels=1,           # le nombre de canaux de sortie
        layers_per_block=2,       # le nombre de couches ResNet à utiliser par bloc UNet
        block_out_channels=(32, 64, 64), 
        down_block_types=( 
            "DownBlock2D",        # un bloc de sous-échantillonnage ResNet standard
            "AttnDownBlock2D",    # un bloc de sous-échantillonnage ResNet avec auto-attention spatiale
            "AttnDownBlock2D",
        ), 
        up_block_types=(
            "AttnUpBlock2D", 
            "AttnUpBlock2D",      # un bloc de suréchantillonnage ResNet avec auto-attention spatiale
            "UpBlock2D",          # un bloc de suréchantillonnage ResNet standard
          ),
    )

  # Notre méthode de transfert prend maintenant les étiquettes de la classe comme argument supplémentaire
  def forward(self, x, t, class_labels):
    # Forme de x :
    bs, ch, w, h = x.shape
    
    # conditionnement de la classe en bon état pour l'ajouter comme canaux d'entrée supplémentaires
    class_cond = self.class_emb(class_labels) # Map to embedding dinemsion
    class_cond = class_cond.view(bs, class_cond.shape[1], 1, 1).expand(bs, class_cond.shape[1], w, h)
    # x est de forme (bs, 1, 28, 28) et class_cond est maintenant (bs, 4, 28, 28)

    # L'entrée nette est maintenant x et la classe cond concaténée ensemble le long de la dimension 1
    net_input = torch.cat((x, class_cond), 1) # (bs, 5, 28, 28)

    # Cette information est transmise à l'UNet en même temps que le pas de temps et renvoie la prédiction
    return self.model(net_input, t).sample # (bs, 1, 28, 28)

Si l’une des formes ou des transformations vous semble confuse, ajoutez des print() pour afficher les formes pertinentes et vérifiez qu’elles correspondent à vos attentes. Nous avons également annoté les formes de certaines variables intermédiaires dans l’espoir de rendre les choses plus claires.

Entraînement et échantillonnage

Alors qu’auparavant nous faisions quelque chose comme prediction = UNet(x, t), nous allons maintenant ajouter les bonnes étiquettes comme troisième argument (prediction = UNet(x, t, y)) pendant l’entraînement, et lors de l’inférence nous pouvons passer les étiquettes que nous voulons et si tout va bien le modèle devrait générer des images qui correspondent. $y$ dans ce cas est l’étiquette des chiffres MNIST, avec des valeurs de 0 à 9.

La boucle d’entraînement est très similaire à l’exemple de l’unité 1. Nous prédisons maintenant le bruit (plutôt que l’image débruitée comme dans l’unité 1) pour correspondre à l’objectif attendu par le DDPMScheduler par défaut que nous utilisons pour ajouter du bruit pendant l’entraînement et pour générer des échantillons au moment de l’inférence. L’entraînement prend du temps. L’accélérer pourrait être un mini-projet amusant, mais la plupart d’entre vous peuvent probablement parcourir le code (et en fait tout ce notebook) sans l’exécuter puisque nous ne faisons qu’illustrer une idée.

# Créer un planificateur
noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='squaredcos_cap_v2')

# Redéfinition du chargeur de données pour fixer la taille du batch à un niveau supérieur à la démonstration de 8
train_dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

# Combien de fois devrions-nous passer les données en revue ?
n_epochs = 10

# Notre réseau 
net = ClassConditionedUNet().to(device)

# Notre fonction de perte
loss_fn = nn.MSELoss()

# L'optimiseur
opt = torch.optim.Adam(net.parameters(), lr=1e-3) 

# Conserver une trace des pertes pour les consulter ultérieurement
losses = []

# La boucle d'entraînement
for epoch in range(n_epochs):
    for x, y in tqdm(train_dataloader):
        
        # Obtenir des données et préparer la version corrompue
        x = x.to(device) * 2 - 1 # Données sur le GPU (sur (-1, 1))
        y = y.to(device)
        noise = torch.randn_like(x)
        timesteps = torch.randint(0, 999, (x.shape[0],)).long().to(device)
        noisy_x = noise_scheduler.add_noise(x, noise, timesteps)

        # Obtenir la prédiction du modèle
        pred = net(noisy_x, timesteps, y) # Notez que nous passons les étiquettes y

        # Calculer la perte
        loss = loss_fn(pred, noise) # Quelle est la distance entre la sortie et le bruit ?

        # Rétropopagation et mise à jour des paramètres :
        opt.zero_grad()
        loss.backward()
        opt.step()

        # Stocker la perte pour plus tard
        losses.append(loss.item())

    # Afficher la moyenne des 100 dernières valeurs de perte pour vous faire une idée de la progression :
    avg_loss = sum(losses[-100:])/100
    print(f'Finished epoch {epoch}. Average of the last 100 loss values: {avg_loss:05f}')

# Visualiser la courbe des pertes
plt.plot(losses)

Une fois l’entraînement terminé, nous pouvons échantillonner quelques images en introduisant différentes étiquettes comme conditionnement :

# Préparer un x aléatoire comme point de départ, ainsi que les étiquettes souhaitées y
x = torch.randn(80, 1, 28, 28).to(device)
y = torch.tensor([[i]*8 for i in range(10)]).flatten().to(device)

# Boucle d'échantillonnage
for i, t in tqdm(enumerate(noise_scheduler.timesteps)):

    # Obtenir la prédiction du modèle
    with torch.no_grad():
        residual = net(x, t, y)  # Notez à nouveau que nous transmettons nos étiquettes y

    # Mise à jour de l'échantillon avec l'étape
    x = noise_scheduler.step(residual, t, x).prev_sample

# Montrer les résultats
fig, ax = plt.subplots(1, 1, figsize=(12, 12))
ax.imshow(torchvision.utils.make_grid(x.detach().cpu().clip(-1, 1), nrow=8)[0], cmap='Greys')

Nous y voilà ! Nous pouvons maintenant contrôler les images produites.

Nous espérons que cet exemple vous a plu. Comme toujours, n’hésitez pas à poser des questions sur Discord.

✏️ À votre tour !
Essayez de refaire la même chose avec FashionMNIST. Modifiez le taux d’apprentissage, la taille du batch et le nombre d’époques.

Pouvez-vous obtenir des images de mode décentes avec moins de temps d’entraînement que dans l’exemple ci-dessus ?

3.1. Introduction à Stable Diffusion

Ce notebook va couvrir les bases de l’utilisation de Stable Diffusion pour créer et modifier des images en utilisant les pipelines existants. Nous allons également jeter un bref coup d’œil aux composants clés au sein du pipeline, tout en laissant une exploration plus approfondie de ces derniers au notebook de plongée profonde. Plus précisément, nous aborderons les points suivants :

• Générer des images à partir d’un texte en utilisant le StableDiffusionPipeline et en expérimentant avec les arguments disponibles
• Voir certains des composants clés du pipeline en action
  • Le VAE qui en fait un « modèle de diffusion latent ».
  • Le tokenizer et l’encodeur qui traitent le prompt textuel
  • L’UNet lui-même
  • Le planificateur et l’exploration de différents planificateurs
• Reproduction de la boucle d’échantillonnage avec les composants du pipeline
• Édition d’images existantes avec le pipeline Img2Img
• Utilisation des pipelines de complétion d’image et Depth2Img

Configuration

!pip install -Uq diffusers ftfy accelerate

# Installer transformers à partir de la source car nous avons besoin de la dernière version pour Depth2Img
!pip install -Uq git+https://github.com/huggingface/transformers 

import torch
import requests
from PIL import Image
from io import BytesIO
from matplotlib import pyplot as plt

# Nous allons explorer un certain nombre de pipelines aujourd'hui !
from diffusers import (
    StableDiffusionPipeline, 
    StableDiffusionImg2ImgPipeline,
    StableDiffusionInpaintPipeline, 
    StableDiffusionDepth2ImgPipeline
    )       

# Nous utiliserons quelques images de démonstration plus loin dans le notebook
def download_image(url):
    response = requests.get(url)
    return Image.open(BytesIO(response.content)).convert("RGB")

# Télécharger des images pour l'exemple de complétion d'image
img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

init_image = download_image(img_url).resize((512, 512))
mask_image = download_image(mask_url).resize((512, 512))

# Définir l'appareil
device = (
    "mps"
    if torch.backends.mps.is_available()
    else "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

Générer des images à partir d’un texte

Chargeons un pipeline Stable Diffusion et voyons ce qu’il peut faire. Il existe plusieurs versions de Stable Diffusion, la plus récente étant la version 2.1. Si vous souhaitez explorer une version plus ancienne, remplacez simplement l’ID du modèle par le modèle approprié (par exemple, vous pouvez essayer “CompVis/stable-diffusion-v1-4” ou choisir un modèle à partir de la bibliothèque de concepts dreambooth).

# Charger le pipeline
model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)

Si vous manquez de mémoire GPU, vous pouvez faire certaines choses pour réduire l’utilisation de la RAM :

• Chargez la version FP16 (non supportée par tous les systèmes). Avec cette version, vous devrez peut-être convertir les tenseurs en torch.float16 lorsque vous expérimenterez avec les composants individuels du pipeline :

  pipe = StableDiffusionPipeline.from_pretrained(model_id, revision="fp16", torch_dtype=torch.float16).to(device)
  

• Activer le découpage de l’attention. Cela permet de réduire l’utilisation de la mémoire du GPU au prix d’une légère réduction de la vitesse :

  pipe.enable_attention_slicing()
  

• Réduire la taille des images générées

Une fois le pipeline chargé, nous pouvons générer une image sur la base d’un prompt avec le code suivant :

# Mise en place d'un générateur pour la reproductibilité
generator = torch.Generator(device=device).manual_seed(42)

# Exécuter le pipeline, en montrant certains des arguments disponibles
pipe_output = pipe(
    prompt="Palette knife painting of an autumn cityscape", # Ce qu'il faut générer
    negative_prompt="Oversaturated, blurry, low quality", # Ce qu'il ne faut PAS générer
    height=480, width=640,     # Spécifier la taille de l'image
    guidance_scale=8,          # Comment suivre avec fermeté le prompt
    num_inference_steps=35,    # Nombre d'étapes
    generator=generator        # Graine aléatoire fixe
)

# Visualiser l'image obtenue :
pipe_output.images[0]

✏️ À votre tour !
Passez un peu de temps à jouer avec la cellule ci-dessus en utilisant vos propres prompts et en modifiant les paramètres pour voir comment ils affectent la sortie. Utilisez une graine aléatoire différente ou supprimez l’argument du générateur pour obtenir des résultats différents à chaque fois. Arguments clés à modifier :

• La largeur et la hauteur spécifient la taille de l’image générée. Elles doivent être divisibles par 8 pour que le VAE fonctionne (ce que nous verrons dans une prochaine section).
• Le nombre de pas influence la qualité de la génération. La valeur par défaut (50) fonctionne bien, mais dans certains cas, on peut se contenter de 20 pas, ce qui est pratique pour l’expérimentation.
• Le prompt négatif est utilisé pendant le processus d’orientation sans classifieur et peut être un moyen utile d’ajouter un contrôle supplémentaire. Vous pouvez ne pas l’utiliser, mais de nombreux utilisateurs trouvent utile de lister certaines descriptions indésirables dans le prompt négatif, comme illustré ci-dessus.
• L’argument guidance_scale détermine l’intensité du guidage sans classifieur (CFG pour classifier-free guidance). Des échelles plus élevées poussent les images générées à mieux correspondre au prompt, mais si l’échelle est trop élevée, les résultats peuvent devenir sursaturés et désagréables.

Si vous souhaitez vous inspirer d’un prompt, le Stable Diffusion Prompt Book est un bon point de départ.

Vous pouvez voir l’effet de l’augmentation de l’échelle d’orientation dans la cellule suivante :

cfg_scales = [1.1, 8, 12] #@param
prompt = "A collie with a pink hat" #@param
fig, axs = plt.subplots(1, len(cfg_scales), figsize=(16, 5))
for i, ax in enumerate(axs):
  im = pipe(prompt, height=480, width=480,
    guidance_scale=cfg_scales[i], num_inference_steps=35,
    generator=torch.Generator(device=device).manual_seed(42)).images[0]
  ax.imshow(im); ax.set_title(f'CFG Scale {cfg_scales[i]}');

Modifiez les valeurs ci-dessus pour essayer différentes échelles et différents prompts. L’interprétation est bien sûr subjective, mais par expérience, toute valeur comprise entre 8 et 12 donne de meilleurs résultats que les valeurs inférieures ou supérieures à cette fourchette.

Composants du pipeline

Le StableDiffusionPipeline que nous utilisons est un peu plus complexe que le DDPMPipeline que nous avons exploré dans les unités précédentes. En plus du UNet et du planificateur, il y a un certain nombre d’autres composants inclus dans le pipeline :

print(list(pipe.components.keys())) # Liste des composants

['vae', 'text_encoder', 'tokenizer', 'unet', 'scheduler', 'safety_checker', 'feature_extractor']

Pour mieux comprendre le fonctionnement du pipeline, voyons brièvement chaque composant en action individuellement, puis assemblons-les pour reproduire la fonctionnalité du pipeline.

Le VAE

Le VAE (auto-encodeur variationnel) est un type de modèle capable d’encoder son entrée dans une représentation comprimée, puis de décoder cette représentation latente pour la rendre proche de l’entrée d’origine. Lors de la génération d’images avec diffusion stable, nous générons d’abord les latents en appliquant le processus de diffusion dans l’espace latent du VAE, puis nous les décodons à la fin pour visualiser l’image résultante.

Voici un code qui prend une image d’entrée, l’encode dans une représentation latente et la décode à nouveau à l’aide de la VAE :

# Créez de fausses données (une image aléatoire, une plage (-1, 1))
images = torch.rand(1, 3, 512, 512).to(device) * 2 - 1 
print("Input images shape:", images.shape)

# Encoder dans l'espace latent
with torch.no_grad():
  latents = 0.18215 * pipe.vae.encode(images).latent_dist.mean
print("Encoded latents shape:", latents.shape)

# Décoder à nouveau
with torch.no_grad():
  decoded_images = pipe.vae.decode(latents / 0.18215).sample
print("Decoded images shape:", decoded_images.shape)

Input images shape: torch.Size([1, 3, 512, 512])
Encoded latents shape: torch.Size([1, 4, 64, 64])
Decoded images shape: torch.Size([1, 3, 512, 512])

Comme vous pouvez le constater, l’image 512x512 est compressée en une représentation latente 64x64 (avec quatre canaux). Cette réduction par 8 de chaque dimension spatiale est la raison pour laquelle la largeur et la hauteur spécifiées doivent être des multiples de 8.

Travailler avec ces latents 4x64x64 riches en informations est plus efficace que de travailler avec des images massives de 512 px, ce qui permet d’obtenir des modèles de diffusion plus rapides dont l’entraînement et l’utilisation nécessitent moins de ressources. Le processus de décodage du VAE n’est pas parfait, mais il est suffisamment bon pour que le petit compromis de qualité en vaille généralement la peine.

NB : L’exemple de code ci-dessus inclut un facteur d’échelle de 0,18215 nécessaire pour correspondre au traitement utilisé lors de l’entraînement de SD.

Le tokenizer et l’encodeur

L’objectif de l’encodeur est de transformer une chaîne d’entrée (le prompt) en une représentation numérique qui peut être transmise à l’UNet en tant que conditionnement. Le texte est d’abord transformé en une série de tokens à l’aide du tokenizer du pipeline. L’encodeur dispose d’un vocabulaire d’environ 50K tokens et tout mot ne figurant pas dans ce vocabulaire est divisé en sous-mots plus petits. Les tokens sont ensuite transmis à l’encodeur lui-même : un transformer qui a été entraîné à l’origine comme encodeur pour CLIP. Nous espérons que ce transformer pré-entraîné a appris des représentations riches du texte qui seront également utiles pour la tâche de diffusion.

Testons ce processus en encodant un prompt d’exemple, d’abord en le tokenizant manuellement et en le faisant passer par l’encodeur puis en utilisant la méthode _encode_prompt pour montrer le processus complet, y compris le rembourrage/la troncature de la longueur jusqu’à la longueur maximale de 77 tokens :

# Tokenizer et encoder un prompt d'exemple manuellement :

# Tokenizer
input_ids = pipe.tokenizer(["A painting of a flooble"])['input_ids']
print("Input ID -> decoded token")
for input_id in input_ids[0]:
  print(f"{input_id} -> {pipe.tokenizer.decode(input_id)}")

# Passage par l'encodeur de texte CLIP
input_ids = torch.tensor(input_ids).to(device)
with torch.no_grad():
  text_embeddings = pipe.text_encoder(input_ids)['last_hidden_state']
print("Text embeddings shape:", text_embeddings.shape)

Input ID -> decoded token
49406 -> <|startoftext|>
320 -> a
3086 -> painting
539 -> of
320 -> a
4062 -> floo
1059 -> ble
49407 -> <|endoftext|>
Text embeddings shape: torch.Size([1, 8, 1024])

# Obtenir les enchâssements finaux à l'aide de la fonction _encode_prompt du pipeline :
text_embeddings = pipe._encode_prompt("A painting of a flooble", device, 1, False, '')
text_embeddings.shape

torch.Size([1, 77, 1024])

Ces enchâssements (les “états cachés” du dernier bloc de transformation dans le modèle de l’encodeur) seront transmis à l’UNet en tant qu’argument supplémentaire de la méthode forward, que nous verrons dans la section suivante.

L’UNet

L’UNet prend une entrée bruitée et prédit le bruit, tout comme les UNets que nous avons vus dans les unités précédentes. Contrairement aux exemples précédents, l’entrée n’est pas une image mais une représentation latente d’une image. En plus du conditionnement du pas de temps, ce UNet prend également en compte l’enchâssement du prompt en tant qu’entrée supplémentaire. Ici, il fait des prédictions sur des données fictives :

# Entrées fictives
timestep = pipe.scheduler.timesteps[0]
latents = torch.randn(1, 4, 64, 64).to(device)
text_embeddings = torch.randn(1, 77, 1024).to(device)

# Prédiction du modèle
with torch.no_grad():
  unet_output = pipe.unet(latents, timestep, text_embeddings).sample
print('UNet output shape:', unet_output.shape) # Même forme que les latents d'entrée

UNet output shape: torch.Size([1, 4, 64, 64])

Le planificateur

Le planificateur stocke le plan de bruit et gère la mise à jour de l’échantillon bruité sur la base des prédictions du modèle. Le planificateur par défaut est un PNDMScheduler, mais vous pouvez en utiliser d’autres (comme LMSDiscreteScheduler) tant qu’ils sont initialisés avec la même configuration.

Nous pouvons tracer le plan de bruit pour voir le niveau de bruit (basé sur $bar{\alpha}$) au fil du temps :

plt.plot(pipe.scheduler.alphas_cumprod, label=r'$\bar{\alpha}$')
plt.xlabel('Timestep (high noise to low noise ->)')
plt.title('Noise schedule');plt.legend()

Si vous souhaitez essayer un autre planificateur, vous pouvez le remplacer comme suit :

from diffusers import LMSDiscreteScheduler

# Remplacer le planificateur
pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)

# Afficher la configuration
print('Scheduler config:', pipe.scheduler)

# Générer une image avec ce nouveau planificateur
pipe(prompt="Palette knife painting of an winter cityscape", height=480, width=480,
     generator=torch.Generator(device=device).manual_seed(42)).images[0]

Scheduler config: LMSDiscreteScheduler {
  "_class_name": "LMSDiscreteScheduler",
  "_diffusers_version": "0.11.1",
  "beta_end": 0.012,
  "beta_schedule": "scaled_linear",
  "beta_start": 0.00085,
  "clip_sample": false,
  "num_train_timesteps": 1000,
  "prediction_type": "epsilon",
  "set_alpha_to_one": false,
  "skip_prk_steps": true,
  "steps_offset": 1,
  "trained_betas": null
}

Vous pouvez lire plus de détails sur l’utilisation de différents planificateurs ici.

Une boucle d’échantillonnage par vous-même

Maintenant que nous avons vu tous ces composants en action, nous pouvons les assembler pour reproduire la fonctionnalité du pipeline :

guidance_scale = 8 #@param
num_inference_steps=30 #@param
prompt = "Beautiful picture of a wave breaking" #@param
negative_prompt = "zoomed in, blurry, oversaturated, warped" #@param

# Encoder le prompt
text_embeddings = pipe._encode_prompt(prompt, device, 1, True, negative_prompt)

# Créer notre point de départ aléatoire
latents = torch.randn((1, 4, 64, 64), device=device, generator=generator)
latents *= pipe.scheduler.init_noise_sigma

# Preparer le planificateur
pipe.scheduler.set_timesteps(num_inference_steps, device=device)

# Boucle sur les pas de temps d'échantillonnage
for i, t in enumerate(pipe.scheduler.timesteps):

  # développer les latents si l'on procède à un guidage sans classifieur
  latent_model_input = torch.cat([latents] * 2)

  # Appliquer tout échelonnement requis par le planificateur
  latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)

  # prédire le bruit résiduel avec l'UNet
  with torch.no_grad():
    noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

  # réaliser un guidage
  noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
  noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

  # calculer l'échantillon bruité précédent x_t -> x_t-1
  latents = pipe.scheduler.step(noise_pred, t, latents).prev_sample

# Décoder les latents résultants en une image
with torch.no_grad():
  image = pipe.decode_latents(latents.detach())

# Visualiser
pipe.numpy_to_pil(image)[0]

Dans la plupart des cas, il sera plus facile d’utiliser les pipelines existants, mais le fait de disposer de cette boucle d’échantillonnage bidouillable peut s’avérer utile pour comprendre et modifier le fonctionnement de chaque composant. Si vous souhaitez voir ce code et tous les différents composants explorés et modifiés en profondeur, consultez le notebook et la video ‘Stable Diffusion Deep Dive’ pour une exploration plus approfondie.

Pipelines supplémentaires

Que pouvons-nous faire d’autre que de générer des images à partir d’un prompt ? Beaucoup de choses ! Dans cette section, nous allons démontrer quelques pipelines sympas pour vous donner un avant-goût des autres tâches pour lesquelles Stable Diffusion peut être utilisé. Plusieurs d’entre eux nécessitent le téléchargement de nouveaux modèles, donc si vous êtes pressé, vous pouvez parcourir cette section en vous contentant de regarder les résultats existants plutôt que de télécharger et d’exécuter tous les modèles vous-même.

Img2Img

Dans les exemples présentés jusqu’à présent, nous avons généré des images en partant de latents aléatoires et en appliquant la boucle complète d’échantillonnage par diffusion. Mais il n’est pas nécessaire de partir de zéro. Le pipeline Img2Img encode d’abord une image existante dans un ensemble de latents, puis ajoute du bruit aux latents et utilise cette image comme point de départ. La quantité de bruit ajoutée et le nombre d’étapes de débruitage appliquées déterminent la « force » du processus Img2Img. L’ajout d’une petite quantité de bruit (force faible) n’entraînera que très peu de changements, tandis que l’ajout d’une quantité maximale de bruit et l’exécution du processus de débruitage complet donneront une image qui ne ressemblera guère à l’image d’entrée, hormis quelques similitudes au niveau de la structure générale.

Ce pipeline ne nécessite pas de modèles particuliers, et tant que l’ID du modèle est le même que celui de notre exemple texte-image ci-dessus, aucun nouveau fichier ne devra être téléchargé.

# Chargement d'un pipeline Img2Img
model_id = "stabilityai/stable-diffusion-2-1-base"
img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id).to(device)

Dans la section Configuration nous avons chargé un exemple init_image à utiliser pour cette démo, mais vous pouvez le remplacer par votre propre image si vous préférez. Voici le pipeline en action :

# Appliquer Img2Img
result_image = img2img_pipe(
    prompt="An oil painting of a man on a bench",
    image = init_image, # L'image de départ
    strength = 0.6, # 0 pour aucun changement, 1.0 pour une force maximale
).images[0]

# Voir le résultat
fig, axs = plt.subplots(1, 2, figsize=(12, 5))
axs[0].imshow(init_image);axs[0].set_title('Input Image')
axs[1].imshow(result_image);axs[1].set_title('Result')

✏️ À votre tour !
Expérimentez avec ce pipeline. Essayez vos propres images, ou jouez avec différentes forces et différents prompts. Vous pouvez utiliser les mêmes arguments que pour le pipeline texte-image, alors n’hésitez pas à essayer différentes tailles, différents nombres d’étapes, etc.

Complétion d’image (inpainting)

Que se passerait-il si nous voulions conserver une partie de l’image d’entrée inchangée mais générer quelque chose de nouveau dans d’autres parties ? C’est ce qu’on appelle la complétion d’image (inpainting). Bien qu’il soit possible de le faire avec le même modèle que les démonstrations précédentes (via StableDiffusionInpaintPipelineLegacy), nous pouvons obtenir de meilleurs résultats en utilisant une version finetunée personnalisée de Stable Diffusion qui prend un masque comme condition supplémentaire. L’image du masque doit avoir la même forme que l’image d’entrée, avec du blanc dans les zones à remplacer et du noir dans les zones à garder inchangées. Voici comment charger un tel pipeline et l’appliquer à l’image d’exemple et au masque chargés dans la section Configuration :

# Charger le pipeline de complétion d'image (nécessite un modèle de complétion d'image approprié)
pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")
pipe = pipe.to(device)

# Complétion d'image avec un prompt pour avoir le résultat souhaité
prompt = "A small robot, high resolution, sitting on a park bench"
image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]

# Voir le résultat
fig, axs = plt.subplots(1, 3, figsize=(16, 5))
axs[0].imshow(init_image);axs[0].set_title('Input Image')
axs[1].imshow(mask_image);axs[1].set_title('Mask')
axs[2].imshow(image);axs[2].set_title('Result')

Ce modèle peut être particulièrement puissant lorsqu’il est combiné à un autre modèle pour générer automatiquement des masques. Par exemple, ce Space utilise un modèle appelé CLIPSeg pour masquer un objet à remplacer sur la base d’une description textuelle.

En marge : gestion du cache de votre modèle

L’exploration de différents pipelines et variantes de modèles peut remplir votre espace disque. Vous pouvez voir quels modèles sont actuellement téléchargés avec :

!ls ~/.cache/huggingface/diffusers/ # Liste du contenu du répertoire cache

Consultez la documentation sur la mise en cache pour savoir comment visualiser et gérer efficacement votre cache.

Depth2Image

Image d’entrée, image de profondeur et exemples générés (source de l’image : StabilityAI)

Img2Img est très bien, mais parfois nous voulons créer une nouvelle image avec la composition de l’original mais avec des couleurs ou des textures complètement différentes. Il peut être difficile de trouver une force d’Img2Img qui préserve ce que nous voulons de la mise en page sans conserver les couleurs d’entrée.

Il est temps d’adopter un autre modèle finetuné ! Celui-ci prend en compte les informations de profondeur comme condition supplémentaire lors de la génération. Le pipeline utilise un modèle d’estimation de la profondeur pour créer une carte de profondeur, qui est ensuite transmise au UNet finetuné lors de la génération d’images afin de préserver (si possible) la profondeur et la structure de l’image initiale tout en remplissant un contenu complètement nouveau.

# Charger le pipeline Depth2Img (nécessite un modèle approprié)
pipe = StableDiffusionDepth2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-depth")
pipe = pipe.to(device)

# Complétion d'image avec un prompt pour avoir le résultat souhaité
prompt = "An oil painting of a man on a bench"
image = pipe(prompt=prompt, image=init_image).images[0]

# Voir le résultat
fig, axs = plt.subplots(1, 2, figsize=(16, 5))
axs[0].imshow(init_image);axs[0].set_title('Input Image')
axs[1].imshow(image);axs[1].set_title('Result');

Notez la comparaison avec l’exemple img2img. Ici, la variation de couleur est beaucoup plus importante mais la structure globale reste fidèle à l’original. Ce n’est pas idéal dans ce cas, car l’homme a été doté d’une anatomie extrêmement bizarre pour correspondre à la forme du chien, mais dans certains cas, c’est extraordinairement utile. Pour un exemple de cette approche, regardez ce tweet montrant le modèle de profondeur utilisé pour texturer une scène en 3D !

Et maintenant ?

Nous espérons vous avoir donné un avant-goût des nombreuses possibilités offertes par Stable Diffusion ! Une fois que vous en aurez assez de jouer avec les exemples de ce notebook, allez voir le notebook du hackathon DreamBooth pour voir comment finetuner votre propre version de Stable Diffusion qui peut être utilisée avec les pipelines texte-image ou img2img que nous avons vus ici.

Si vous êtes curieux d’approfondir le fonctionnement des différents composants, consultez le notebook Stable Diffusion Deep Dive qui va beaucoup plus loin dans les détails et montre quelques astuces supplémentaires que nous pouvons faire.

N’oubliez pas de partager vos créations avec nous et la communauté !

3.2. Stable Diffusion : plongée en profondeur

Stable Diffusion est un puissant modèle de texte à image. Il existe plusieurs sites web et outils pour rendre son utilisation aussi simple que possible. Il est également intégré à la bibliothèque de Diffusers d’Huggingface, ce qui permet de générer des images en toute simplicité :

from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=True).to("cuda")
image = pipe("An astronaught scuba diving").images[0]

Dans ce notebook, nous allons nous plonger dans le code qui se cache derrière ces interfaces faciles à utiliser, pour voir ce qui se passe sous le capot. Nous commencerons par recréer la fonctionnalité ci-dessus sous la forme d’un morceau de code effrayant, puis, un par un, nous inspecterons les différents composants et comprendrons ce qu’ils font. À la fin de ce notebook, cette même boucle d’échantillonnage devrait ressembler à quelque chose que vous pouvez peaufiner et modifier à votre guise.

Configuration et importations

Vous devrez vous connecter à Hugging Face et accepter les termes de la licence pour ce modèle (voir la carte de modèle pour plus de détails). Lorsque vous exécuterez ce notebook pour la première fois, vous devrez décommenter les deux cellules suivantes pour installer les prérequis et vous connecter au Hub avec un token d’accès.

# !pip install -q --upgrade transformers diffusers ftfy

from base64 import b64encode

import numpy
import torch
from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
from huggingface_hub import notebook_login

# Pour l'affichage vidéo
from IPython.display import HTML
from matplotlib import pyplot as plt
from pathlib import Path
from PIL import Image
from torch import autocast
from torchvision import transforms as tfms
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer, logging

torch.manual_seed(1)
if not (Path.home()/'.huggingface'/'token').exists(): notebook_login()

# Suppression de certains avertissements inutiles lors du chargement de CLIPTextModel
logging.set_verbosity_error()

# Définir l'appareil
torch_device = "cuda" if torch.cuda.is_available() else "cpu"

Chargement des modèles

Ce code (et celui de la section suivante) provient du notebook illustratif d’Huggingface.

Il télécharge et configure les modèles et les composants que nous utiliserons. Exécutons-le pour l’instant et passons à la section suivante pour vérifier que tout fonctionne avant d’aller plus loin.

Si vous avez chargé un pipeline, vous pouvez aussi accéder à ces composants en utilisant pipe.unet, pipe.vae et ainsi de suite.

Dans ce notebook, nous ne faisons pas d’économies de mémoire. Si vous vous retrouvez à court de RAM GPU, regardez le code du pipeline pour vous inspirer avec des choses comme le découpage de l’attention, le passage à la demi-précision (fp16), le maintien du VAE sur le CPU et d’autres modifications.

# Charger le modèle auto-encodeur qui sera utilisé pour décoder les latents dans l'espace de l'image
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")

# Charger le tokenizer et l'encodeur
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")

# Le modèle UNet pour générer les latents
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")

# Le planificateur de bruit
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)

# Nous allons au GPU !
vae = vae.to(torch_device)
text_encoder = text_encoder.to(torch_device)
unet = unet.to(torch_device)

Une boucle de diffusion

Si tout ce que vous voulez, c’est créer une image avec du texte, vous pouvez ignorer ce notebook et utiliser l’un des outils existants (comme DreamStudio) ou utiliser le pipeline simplifié d’Hugging Face comme documenté ici.

Ce que nous voulons faire ici, c’est approfondir un peu plus la façon dont cela fonctionne. Nous allons donc commencer par vérifier que le code de l’exemple s’exécute. Il ressemble beaucoup à ce que vous trouverez si vous inspectez la méthode call() du pipeline de Stable Diffusion.

# Quelques paramètres
prompt = ["A watercolor painting of an otter"]
height = 512                        # hauteur par défaut de Stable Diffusion
width = 512                         # largeur par défaut de Stable Diffusion
num_inference_steps = 30            # Nombre d'étapes de débruitage
guidance_scale = 7.5                # Échelle pour un guidage sans classifieur
generator = torch.manual_seed(32)   # Générateur de la graine pour créer le bruit latent initial
batch_size = 1

# Preparation du texte
text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
    text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
    [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
with torch.no_grad():
    uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0] 
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

# Preparation du planificateur
scheduler.set_timesteps(num_inference_steps)

# Preparation des latents
latents = torch.randn(
  (batch_size, unet.in_channels, height // 8, width // 8),
  generator=generator,
)
latents = latents.to(torch_device)
latents = latents * scheduler.init_noise_sigma # Mise à l'échelle (versions précédentes) latents = latents * self.scheduler.sigmas[0]

# Boucle
with autocast("cuda"):
    for i, t in tqdm(enumerate(scheduler.timesteps)):
        # étendre les latents si nous procédons à un guidage sans classifieur afin d'éviter de faire deux passages en avant        
		latent_model_input = torch.cat([latents] * 2)
        sigma = scheduler.sigmas[i]
        # mettre à l'échelle les latents (préconditionnement)
        # latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5) # Diffusers 0.3 et moins
        latent_model_input = scheduler.scale_model_input(latent_model_input, t)

        # prédire le bruit résiduel
        with torch.no_grad():
            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

        # effectuer le guidage
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

        # calculer l'échantillon bruité précédent x_t -> x_t-1
        # latents = scheduler.step(noise_pred, i, latents)["prev_sample"] # Diffusers 0.3 et moins
        latents = scheduler.step(noise_pred, t, latents).prev_sample

# mettre à l'échelle et décoder les latents de l'image à l'aide du vae
latents = 1 / 0.18215 * latents
with torch.no_grad():
    image = vae.decode(latents).sample

# Affichage
image = (image / 2 + 0.5).clamp(0, 1)
image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
images = (image * 255).round().astype("uint8")
pil_images = [Image.fromarray(image) for image in images]
pil_images[0]

Cela fonctionne, mais cela fait beaucoup de code ! Examinons les composants un par un.

L’auto-encodeur (AE)

L’AE peut encoder une image dans une sorte de représentation latente, et la décoder à nouveau en une image. Nous avons regroupé le code dans quelques fonctions pour que nous puissions voir à quoi cela ressemble en action :

def pil_to_latent(input_im):
    # Une seule image -> un seul latent dans un batch (donc taille 1, 4, 64, 64)
    with torch.no_grad():
        latent = vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling
    return 0.18215 * latent.latent_dist.sample()

def latents_to_pil(latents):
    # bain de latents -> liste d'images
    latents = (1 / 0.18215) * latents
    with torch.no_grad():
        image = vae.decode(latents).sample
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
    images = (image * 255).round().astype("uint8")
    pil_images = [Image.fromarray(image) for image in images]
    return pil_images

Nous utiliserons ici une image provenant du web, mais vous pouvez charger la vôtre en la téléchargeant et en modifiant le nom du fichier dans la cellule suivante.

# Télécharger une image de démonstration
!curl --output macaw.jpg 'https://lafeber.com/pet-birds/wp-content/uploads/2018/06/Scarlet-Macaw-2.jpg'

  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100 62145  100 62145    0     0  10874      0  0:00:05  0:00:05 --:--:-- 15633

# Charger l'image avec PIL
input_image = Image.open('macaw.jpg').resize((512, 512))
input_image

L’encodage dans l’espace latent de l’AE à l’aide de la fonction définie ci-dessus se présente comme suit :

# Encoder dans l'espace latent
encoded = pil_to_latent(input_image)
encoded.shape

torch.Size([1, 4, 64, 64])

# Visualisons les quatre canaux de cette représentation latente :
fig, axs = plt.subplots(1, 4, figsize=(16, 4))
for c in range(4):
    axs[c].imshow(encoded[0][c].cpu(), cmap='Greys')

Ce tenseur 4x64x64 capture de nombreuses informations sur l’image, suffisamment, espérons-le, pour que lorsque nous l’introduisons dans le décodeur, nous obtenions en retour quelque chose de très proche de notre image d’entrée :

# Décoder cette représentation latente en une image
decoded = latents_to_pil(encoded)[0]
decoded

Vous verrez de petites différences si vous plissez les yeux ! Concentrez-vous sur l’œil si vous ne voyez rien d’évident. C’est assez impressionnant : cette image latente de 4x64x64 semble contenir beaucoup plus d’informations qu’une image de 64px.

Cet auto-encodeur a été entraîné à réduire une image à une représentation plus petite, puis à recréer l’image à partir de cette version compressée.

Dans ce cas particulier, le facteur de compression est de 48, nous partons d’une image 3x512x512(cannaux x hauteur x largeur) et elle est compressée en un vecteur latent 4x64x64. Chaque volume de 3x8x8 pixels dans l’image d’entrée est compressé en seulement 4 nombres (4x1x1). Il est possible de trouver des AEs avec un taux de compression plus élevé (par exemple f16 comme certains modèles populaires de VQGAN) mais à un moment donné, ils commencent à introduire des artefacts que nous ne voulons pas.

Pourquoi utiliser un auto-encodeur ? Nous pouvons faire de la diffusion dans l’espace des pixels où le modèle reçoit toutes les données de l’image comme entrées et produit une prédiction de sortie de la même forme. Mais cela implique le traitement d’un grand nombre de données et rend la génération d’images à haute résolution très coûteuse sur le plan informatique. Certaines solutions consistent à effectuer la diffusion à basse résolution (64 px par exemple), puis à entraîner un modèle distinct pour augmenter l’échelle de manière répétée (comme avec D2/Imagen). La diffusion latente, quant à elle, effectue le processus de diffusion dans cet espace latent, en utilisant les représentations compressées de notre AE plutôt que des images brutes. Ces représentations sont riches en informations et peuvent être suffisamment petites pour être gérées par du matériel grand public. Une fois que nous avons généré une nouvelle image en tant que représentation latente, l’auto-encodeur peut prendre ces sorties latentes finales et les transformer en pixels réels.

Le planificateur

Nous devons maintenant parler de l’ajout de bruit.

Pendant l’entraînement, nous ajoutons du bruit à une image, puis nous demandons au modèle d’essayer de prédire le bruit. Si nous ajoutons toujours beaucoup de bruit, le modèle risque de ne pas avoir grand-chose à faire. Si nous n’en ajoutons qu’une infime quantité, le modèle ne pourra pas faire grand-chose avec les points de départ aléatoires que nous utilisons pour l’échantillonnage. Au cours de l’entraînement, la quantité de bruit varie donc en fonction d’une certaine distribution.

Pendant l’échantillonnage, nous voulons « débruiter » sur un certain nombre d’étapes. Le nombre d’étapes et la quantité de bruit que nous devons viser à chaque étape affecteront le résultat final.

Le planificateur est chargé de gérer tous ces détails. Par exemple : scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000) met en place un scheduler qui correspond à celui utilisé pour entraîner ce modèle. Lorsque nous voulons échantillonner sur un plus petit nombre de pas, nous le faisons avec scheduler.set_timesteps :

# Réglage du nombre de pas d'échantillonnage :
scheduler.set_timesteps(15)

Vous pouvez voir comment notre nouvel ensemble d’étapes correspond à celles utilisées dans l’entraînement :

# Voyez ça en termes de 1000 étapes originales utilisées pour l'entraînement :
print(scheduler.timesteps)

tensor([999.0000, 927.6429, 856.2857, 784.9286, 713.5714, 642.2143, 570.8571,
        499.5000, 428.1429, 356.7857, 285.4286, 214.0714, 142.7143,  71.3571,
          0.0000], dtype=torch.float64)

Et quelle est la quantité de bruit présente à chaque endroit :

# Examinez les niveaux de bruit équivalents :
print(scheduler.sigmas)

tensor([14.6146,  9.6826,  6.6780,  4.7746,  3.5221,  2.6666,  2.0606,  1.6156,
         1.2768,  1.0097,  0.7913,  0.6056,  0.4397,  0.2780,  0.0292,  0.0000])

Pendant l’échantillonnage, nous partons d’un niveau de bruit élevé (en fait, notre entrée sera du bruit pur) et nous « débruitons » progressivement jusqu’à obtenir une image, selon ce calendrier.

# Affichage du planificateur de bruit :
plt.plot(scheduler.sigmas)
plt.title('Noise Schedule')
plt.xlabel('Sampling step')
plt.ylabel('sigma')
plt.show()

Ce « sigma » est la quantité de bruit ajoutée à la représentation latente. Voyons ce que cela donne en ajoutant un peu de bruit à notre image codée, puis en décodant cette version bruitée :

noise = torch.randn_like(encoded) # Bruit aléatoire
sampling_step = 10 # Equivalent à une étape 10 sur 15 dans la grille ci-dessus
# encoded_and_noised = scheduler.add_noise(encoded, noise, timestep) # Diffusers 0.3 et en dessous
encoded_and_noised = scheduler.add_noise(encoded, noise, timesteps=torch.tensor([scheduler.timesteps[sampling_step]]))
latents_to_pil(encoded_and_noised.float())[0] # Affichage

À quoi cela ressemble-t-il à différents pas de temps ? Faites l’expérience et voyez par vous-même !

Si vous décommentez la cellule ci-dessous, vous verrez que dans ce cas, la fonction scheduler.add_noise ne fait qu’ajouter du bruit à l’échelle sigma : noisy_samples = original_samples + noise * sigmas

# ??scheduler.add_noise

D’autres modèles de diffusion peuvent être entraînés avec différentes approches de bruits et d’ordonnancement, dont certaines maintiennent la variance relativement constante entre les niveaux de bruit (« préservation de la variance ») avec différentes astuces de mise à l’échelle et de mélange au lieu d’avoir des latents bruités avec une variance de plus en plus élevée au fur et à mesure que l’on ajoute du bruit (« explosion de la variance »).

Si nous voulons partir d’un bruit aléatoire au lieu d’une image bruitée, nous devons la mettre à l’échelle de la plus grande valeur sigma utilisée pendant l’entraînement, soit ~14 dans ce cas. Et avant que ces latents bruités ne soient introduits dans le modèle, ils sont à nouveau mis à l’échelle dans l’étape dite de pré-conditionnement : latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5) (maintenant géré par latent_model_input = scheduler.scale_model_input(latent_model_input, t)).

Encore une fois, cette mise à l’échelle/pré-conditionnement diffère entre les articles et les implémentations, alors gardez un œil sur ce point si vous travaillez avec un type différent de modèle de diffusion.

Boucle à partir de la version bruitée de l’entrée (AKA image2image)

Voyons ce qui se passe lorsque nous utilisons notre image comme point de départ, en ajoutant un peu de bruit et en effectuant les dernières étapes de débruitage dans la boucle avec un nouveau prompt.

Nous allons utiliser une boucle similaire à celle de la première démonstration, mais nous allons sauter les premières étapes start_step.

Pour bruiter notre image, nous utiliserons un code comme celui montré ci-dessus, en utilisant le planificateur pour la bruiter à un niveau équivalent à l’étape 10 (start_step).

# Paramètres (les mêmes que précédemment, à l'exception du nouveau prompt)
prompt = ["A colorful dancer, nat geo photo"]
height = 512                        # hauteur par défaut de Stable Diffusion
width = 512                         # largeur par défaut de Stable Diffusion
num_inference_steps = 30            # Nombre d'étapes de débruitage
guidance_scale = 7.5                # Échelle pour un guidage sans classifieur
generator = torch.manual_seed(32)   # Générateur de la graine pour créer le bruit latent initial
batch_size = 1

# Preparation du texte (comme précédemment)
text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
    text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
    [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
with torch.no_grad():
    uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0] 
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

# Preparation du planificateur (définition du nombre d'étapes de l'inférence)
scheduler.set_timesteps(num_inference_steps)

# Preparation des latents (bruitage approprié pour start_step)
start_step = 10
start_sigma = scheduler.sigmas[start_step]
noise = torch.randn_like(encoded)
latents = scheduler.add_noise(encoded, noise, timesteps=torch.tensor([scheduler.timesteps[start_step]]))
latents = latents.to(torch_device).float()

# Boucle
for i, t in tqdm(enumerate(scheduler.timesteps)):
    if i >= start_step: # << C'est la seule modification que nous apportons à la boucle.
        
        # étendre les latents si nous procédons à un guidage sans classifieur afin d'éviter de faire deux passages en avant  
        latent_model_input = torch.cat([latents] * 2)
        sigma = scheduler.sigmas[i]
        latent_model_input = scheduler.scale_model_input(latent_model_input, t)

        # prédire le bruit résiduel
        with torch.no_grad():
            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]

        # effectuer le guidage
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

        # calculer l'échantillon bruité précédent x_t -> x_t-1
        latents = scheduler.step(noise_pred, t, latents).prev_sample

latents_to_pil(latents)[0]

Vous pouvez voir que certaines couleurs et structures de l’image sont conservées, mais nous avons maintenant une nouvelle image ! Plus vous ajoutez de bruit et plus vous effectuez d’étapes, plus l’image s’éloigne de l’image d’entrée.

C’est ainsi que fonctionne le célèbre pipeline img2img. Encore une fois, si c’est votre objectif final, il existe des outils qui facilitent la tâche !

Mais vous pouvez voir que sous le capot, c’est la même chose que la boucle de génération, en sautant les premières étapes et en partant d’une image bruitée plutôt que d’une image purement bruitée.

Essayez de changer le nombre d’étapes sautées et de voir comment cela affecte la quantité de changement de l’image par rapport à l’entrée.

Exploration du pipeline texte -> enchâssement

Nous utilisons un modèle d’encodage de texte pour transformer notre texte en un ensemble d’enchâssements qui sont transmis au modèle de diffusion en tant que conditionnement. Suivons un morceau de texte tout au long de ce processus et voyons comment il fonctionne.

# Notre prompt textuel
prompt = 'A picture of a puppy'

Nous commençons par la tokenisation :

# Transformer le texte en une séquence de tokens :
text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
text_input['input_ids'][0] # Voir les tokens

tensor([49406,   320,  1674,   539,   320,  6829, 49407, 49407, 49407, 49407,
        49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
        49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
        49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
        49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
        49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
        49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
        49407, 49407, 49407, 49407, 49407, 49407, 49407])

# Voir les tokens individuels
for t in text_input['input_ids'][0][:8]: # Nous nous contenterons d'examiner les 7 premiers pour vous éviter un mur d'<|endoftext|>'
    print(t, tokenizer.decoder.get(int(t)))

tensor(49406) <|startoftext|>
tensor(320) a</w>
tensor(1674) picture</w>
tensor(539) of</w>
tensor(320) a</w>
tensor(6829) puppy</w>
tensor(49407) <|endoftext|>
tensor(49407) <|endoftext|>

Nous pouvons passer directement aux enchâssements finaux (de sortie) de la manière suivante :

# Récupérer les enchâssements de sortie
output_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
print('Shape:', output_embeddings.shape)
output_embeddings

Shape: torch.Size([1, 77, 768])

tensor([[[-0.3884,  0.0229, -0.0522,  ..., -0.4899, -0.3066,  0.0675],
         [ 0.0290, -1.3258,  0.3085,  ..., -0.5257,  0.9768,  0.6652],
         [ 0.6942,  0.3538,  1.0991,  ..., -1.5716, -1.2643, -0.0121],
         ...,
         [-0.0221, -0.0053, -0.0089,  ..., -0.7303, -1.3830, -0.3011],
         [-0.0062, -0.0246,  0.0065,  ..., -0.7326, -1.3745, -0.2953],
         [-0.0536,  0.0269,  0.0444,  ..., -0.7159, -1.3634, -0.3075]]],
       device='cuda:0', grad_fn=<NativeLayerNormBackward0>)

Nous passons nos tokens à travers text_encoder et nous obtenons comme par magie des nombres que nous pouvons introduire dans le modèle.

Comment ces chiffres sont-ils générés ? Les tokens sont transformés en un ensemble d’enchâssements d’entrée, qui sont ensuite introduits dans le transformer pour obtenir les enchâssements de sortie finaux.

Pour obtenir ces enchâssements d’entrée, il y a en fait deux étapes comme le révèle l’inspection de text_encoder.text_model.embeddings :

text_encoder.text_model.embeddings

CLIPTextEmbeddings(
  (token_embedding): Embedding(49408, 768)
  (position_embedding): Embedding(77, 768)
)

Enchâssement de tokens

Le token est envoyé à la fonction token_embedding pour le transformer en vecteur. Le nom de la fonction get_input_embeddings est trompeur puisque ces enchâssements de tokens doivent être combinés avec les enchâssements de positions avant d’être utilisés comme entrées dans le modèle ! Quoi qu’il en soit, examinons d’abord la partie relative à l’enchâssements des tokens.

Nous pouvons regarder la couche d’enchâssement :

# Accéder à la couche enchâssement
token_emb_layer = text_encoder.text_model.embeddings.token_embedding
token_emb_layer # Taille du vocabulaire 49408, emb_dim 768

Embedding(49408, 768)

Et enchâsser un token comme suit :

# Enchâsser un *token*, dans ce cas, celui du "chiot"
embedding = token_emb_layer(torch.tensor(6829, device=torch_device))
embedding.shape # représentation en 768-dim

torch.Size([768])

Cet unique tokens a été associé avec un vecteur à 768 dimensions.

Nous pouvons faire la même chose avec tous les tokens du prompt pour obtenir tous les enchâssements de tokens :

token_embeddings = token_emb_layer(text_input.input_ids.to(torch_device))
print(token_embeddings.shape) # taille du batch 1, 77 *tokens*, 768 valeurs pour chaque
token_embeddings

torch.Size([1, 77, 768])

tensor([[[ 0.0011,  0.0032,  0.0003,  ..., -0.0018,  0.0003,  0.0019],
         [ 0.0013, -0.0011, -0.0126,  ..., -0.0124,  0.0120,  0.0080],
         [ 0.0235, -0.0118,  0.0110,  ...,  0.0049,  0.0078,  0.0160],
         ...,
         [ 0.0012,  0.0077, -0.0011,  ..., -0.0015,  0.0009,  0.0052],
         [ 0.0012,  0.0077, -0.0011,  ..., -0.0015,  0.0009,  0.0052],
         [ 0.0012,  0.0077, -0.0011,  ..., -0.0015,  0.0009,  0.0052]]],
       device='cuda:0', grad_fn=<EmbeddingBackward0>)

Enchâssements positionnels

Les enchâssements positionnels indiquent au modèle à quel endroit d’une séquence se trouve un token. Tout comme l’enchâssement de * tokens, il s’agit d’un ensemble de paramètres (qui peuvent éventuellement être appris). Mais maintenant, au lieu de traiter ~50k *tokens nous avons juste besoin d’un pour chaque position (77 au total) :

pos_emb_layer = text_encoder.text_model.embeddings.position_embedding
pos_emb_layer

Embedding(77, 768)

Nous pouvons obtenir l’enchâssement positionnel pour chaque position :

position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]
position_embeddings = pos_emb_layer(position_ids)
print(position_embeddings.shape)
position_embeddings

torch.Size([1, 77, 768])

tensor([[[ 0.0016,  0.0020,  0.0002,  ..., -0.0013,  0.0008,  0.0015],
         [ 0.0042,  0.0029,  0.0002,  ...,  0.0010,  0.0015, -0.0012],
         [ 0.0018,  0.0007, -0.0012,  ..., -0.0029, -0.0009,  0.0026],
         ...,
         [ 0.0216,  0.0055, -0.0101,  ..., -0.0065, -0.0029,  0.0037],
         [ 0.0188,  0.0073, -0.0077,  ..., -0.0025, -0.0009,  0.0057],
         [ 0.0330,  0.0281,  0.0289,  ...,  0.0160,  0.0102, -0.0310]]],
       device='cuda:0', grad_fn=<EmbeddingBackward0>)

Combiner les enchâssements de tokens et de positions

Il est temps de combiner les deux. Comment faire ? Il suffit de les additionner ! D’autres approches sont possibles, mais pour ce modèle, c’est ainsi que nous procédons.

En les combinant de cette manière, nous obtenons les enchâssements d’entrée finaux, prêts à être introduits dans le transformer :

# En les combinant, nous obtenons les enchâssements d'entrée finaux
input_embeddings = token_embeddings + position_embeddings
print(input_embeddings.shape)
input_embeddings

torch.Size([1, 77, 768])

tensor([[[ 2.6770e-03,  5.2133e-03,  4.9323e-04,  ..., -3.1321e-03,
           1.0659e-03,  3.4316e-03],
         [ 5.5371e-03,  1.7510e-03, -1.2381e-02,  ..., -1.1410e-02,
           1.3508e-02,  6.8378e-03],
         [ 2.5356e-02, -1.1019e-02,  9.7663e-03,  ...,  1.9460e-03,
           6.8375e-03,  1.8573e-02],
         ...,
         [ 2.2781e-02,  1.3262e-02, -1.1241e-02,  ..., -8.0054e-03,
          -2.0560e-03,  8.9366e-03],
         [ 2.0026e-02,  1.5015e-02, -8.7638e-03,  ..., -4.0313e-03,
           1.8487e-05,  1.0885e-02],
         [ 3.4206e-02,  3.5826e-02,  2.7768e-02,  ...,  1.4465e-02,
           1.1110e-02, -2.5745e-02]]], device='cuda:0', grad_fn=<AddBackward0>)

Nous pouvons vérifier que ces résultats sont les mêmes que ceux obtenus avec text_encoder.text_model.embeddings :

# La procédure suivante combine toutes les étapes ci-dessus (mais ne nous permet pas de les modifier !)
text_encoder.text_model.embeddings(text_input.input_ids.to(torch_device))

tensor([[[ 2.6770e-03,  5.2133e-03,  4.9323e-04,  ..., -3.1321e-03,
           1.0659e-03,  3.4316e-03],
         [ 5.5371e-03,  1.7510e-03, -1.2381e-02,  ..., -1.1410e-02,
           1.3508e-02,  6.8378e-03],
         [ 2.5356e-02, -1.1019e-02,  9.7663e-03,  ...,  1.9460e-03,
           6.8375e-03,  1.8573e-02],
         ...,
         [ 2.2781e-02,  1.3262e-02, -1.1241e-02,  ..., -8.0054e-03,
          -2.0560e-03,  8.9366e-03],
         [ 2.0026e-02,  1.5015e-02, -8.7638e-03,  ..., -4.0313e-03,
           1.8487e-05,  1.0885e-02],
         [ 3.4206e-02,  3.5826e-02,  2.7768e-02,  ...,  1.4465e-02,
           1.1110e-02, -2.5745e-02]]], device='cuda:0', grad_fn=<AddBackward0>)

Passage dans le transformer

Nous voulons modifier les enchâssements d’entrée (en particulier les enchâssements de tokens) avant de les envoyer dans le reste du modèle, mais nous devons d’abord nous assurer que nous savons comment le faire. Nous avons lu le code de la méthode forward du text_encoder, et nous nous sommes basés sur ce code pour la méthode forward du text_model que le text_encoder englobe. Pour l’inspecter vous-même, tapez ??text_encoder.text_model.forward et vous obtiendrez les informations sur la fonction et le code source, une astuce de débogage utile !

Quoi qu’il en soit, nous pouvons copier les bits dont nous avons besoin pour obtenir ce que l’on appelle le « dernier état caché » et ainsi générer nos enchâssements finaux :

def get_output_embeds(input_embeddings):
    # Le modèle de texte de CLIP utilise le masquage causal, c'est pourquoi nous le préparons ici :
    bsz, seq_len = input_embeddings.shape[:2]
    causal_attention_mask = text_encoder.text_model._build_causal_attention_mask(bsz, seq_len, dtype=input_embeddings.dtype)

	# Obtenir les enchâssements de sortie implique d'appeler le modèle en passant output_hidden_states=True 
    # afin qu'il ne renvoie pas uniquement les prédictions finales regroupées :
    encoder_outputs = text_encoder.text_model.encoder(
        inputs_embeds=input_embeddings,
        attention_mask=None, # Nous n'utilisons pas de masque d'attention, cela peut donc être None.
        causal_attention_mask=causal_attention_mask.to(torch_device),
        output_attentions=None,
        output_hidden_states=True, # Nous voulons le résultat des enchâssements et non le résultat final.
        return_dict=None,
    )

    # Seul l'état caché de sortie nous intéresse
    output = encoder_outputs[0]

    # Il existe une normalisation de couche finale par laquelle nous devons passer
    output = text_encoder.text_model.final_layer_norm(output)

    # Et maintenant, elles sont prêtes !
    return output

out_embs_test = get_output_embeds(input_embeddings) # Alimenter le modèle à l'aide de notre nouvelle fonction
print(out_embs_test.shape) # Vérifier la forme de la sortie
out_embs_test # Inspecter la sortie

torch.Size([1, 77, 768])

tensor([[[-0.3884,  0.0229, -0.0522,  ..., -0.4899, -0.3066,  0.0675],
         [ 0.0290, -1.3258,  0.3085,  ..., -0.5257,  0.9768,  0.6652],
         [ 0.6942,  0.3538,  1.0991,  ..., -1.5716, -1.2643, -0.0121],
         ...,
         [-0.0221, -0.0053, -0.0089,  ..., -0.7303, -1.3830, -0.3011],
         [-0.0062, -0.0246,  0.0065,  ..., -0.7326, -1.3745, -0.2953],
         [-0.0536,  0.0269,  0.0444,  ..., -0.7159, -1.3634, -0.3075]]],
       device='cuda:0', grad_fn=<NativeLayerNormBackward0>)

Notez que cela correspond aux output_embeddings que nous avons vu au début. Nous avons trouvé comment diviser cette étape (“obtenir les enchâssements”) en plusieurs sous-étapes prêtes à être modifiées.

Maintenant que nous avons mis en place ce processus, nous pouvons remplacer l’encodage d’entrée d’un token par un nouvel encodage de notre choix, ce qui dans notre cas d’utilisation final, sera quelque chose que nous apprendrons. Pour démontrer le concept, remplaçons l’encodage d’entrée de « puppy » dans le prompt avec lequel nous avons joué avec l’enchâssement du token 2368, obtenons un nouvel ensemble d’enchâssement de sortie basés sur celui-ci et utilisons-les pour générer une image afin de voir ce que nous obtenons :

prompt = 'A picture of a puppy'

# Tokeniser
text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
input_ids = text_input.input_ids.to(torch_device)

# Obtenir les enchâssements des tokens
token_embeddings = token_emb_layer(input_ids)

# Le nouvel enchâssement. Dans ce cas, il s'agit simplement de l'enchâssement d'entrée du token 2368
replacement_token_embedding = text_encoder.get_input_embeddings()(torch.tensor(2368, device=torch_device))

# Insérer ceci dans les enchâssements de token
token_embeddings[0, torch.where(input_ids[0]==6829)] = replacement_token_embedding.to(torch_device)

# Combiner avec le'enchâssement positionnel
input_embeddings = token_embeddings + position_embeddings

# Passage dans le transformer pour obtenir les enchâssements finaux
modified_output_embeddings = get_output_embeds(input_embeddings)

print(modified_output_embeddings.shape)
modified_output_embeddings

torch.Size([1, 77, 768])

tensor([[[-0.3884,  0.0229, -0.0522,  ..., -0.4899, -0.3066,  0.0675],
         [ 0.0290, -1.3258,  0.3085,  ..., -0.5257,  0.9768,  0.6652],
         [ 0.6942,  0.3538,  1.0991,  ..., -1.5716, -1.2643, -0.0121],
         ...,
         [-0.6034, -0.5322,  0.0629,  ..., -0.3964,  0.0877, -0.9558],
         [-0.5936, -0.5407,  0.0731,  ..., -0.3876,  0.0906, -0.9436],
         [-0.6393, -0.4703,  0.1103,  ..., -0.3904,  0.1351, -0.9726]]],
       device='cuda:0', grad_fn=<NativeLayerNormBackward0>)

Les premiers sont identiques, les derniers ne le sont pas. Tout ce qui se trouve à la position du token que nous remplaçons et après sera affecté.

Si tout s’est bien passé, nous devrions voir autre chose qu’un chiot lorsque nous les utiliserons pour générer une image. Et bien sûr, c’est le cas !

# Génération d'une image avec ces enchâssements modifiés

def generate_with_embs(text_embeddings):
	height = 512                        # hauteur par défaut de Stable Diffusion
	width = 512                         # largeur par défaut de Stable Diffusion
	num_inference_steps = 30            # Nombre d'étapes de débruitage
	guidance_scale = 7.5                # Échelle pour un guidage sans classifieur
	generator = torch.manual_seed(32)   # Générateur de la graine pour créer le bruit latent initial
    batch_size = 1

    max_length = text_input.input_ids.shape[-1]
    uncond_input = tokenizer(
      [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
    )
    with torch.no_grad():
        uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0] 
    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

    # Preparation du planificateur
    scheduler.set_timesteps(num_inference_steps)

    # Preparation des latents
    latents = torch.randn(
    (batch_size, unet.in_channels, height // 8, width // 8),
    generator=generator,
    )
    latents = latents.to(torch_device)
    latents = latents * scheduler.init_noise_sigma

    # Boucle
    for i, t in tqdm(enumerate(scheduler.timesteps)):
        # étendre les latents si nous procédons à un guidage sans classifieur afin d'éviter de faire deux passages en avant
        latent_model_input = torch.cat([latents] * 2)
        sigma = scheduler.sigmas[i]
        latent_model_input = scheduler.scale_model_input(latent_model_input, t)

        # prédire le bruit résiduel
        with torch.no_grad():
            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]

        # réaliser un guidage
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

        # calculer l'échantillon bruité précédent x_t -> x_t-1
        latents = scheduler.step(noise_pred, t, latents).prev_sample

    return latents_to_pil(latents)[0]

generate_with_embs(modified_output_embeddings)

Surprise ! Vous savez maintenant ce que signifie le token 2368.

Que pouvons-nous en faire ? Pourquoi nous sommes-nous donné tout ce mal ? Eh bien, nous verrons bientôt un cas d’utilisation plus convaincant, mais en résumé, une fois que nous pouvons accéder aux enchâssements de tokens et les modifier, nous pouvons faire des choses comme les remplacer par autre chose. Dans l’exemple que nous venons de faire, il s’agissait simplement d’un autre enchâssement de tokens du vocabulaire du modèle, ce qui équivaut à une simple modification du prompt. Mais nous pouvons également mélanger les tokens. Par exemple, voici un mi-chiot / mi-mouflette :

# Au cas où vous vous demanderiez comment obtenir le *token* d'un mot, ou l'enchâssement d'un *token* :
prompt = 'skunk'
print('tokenizer(prompt):', tokenizer(prompt))
print('token_emb_layer([token_id]) shape:', token_emb_layer(torch.tensor([8797], device=torch_device)).shape)

tokenizer(prompt): {'input_ids': [49406, 42194, 49407], 'attention_mask': [1, 1, 1]}
token_emb_layer([token_id]) shape: torch.Size([1, 768])

prompt = 'A picture of a puppy'

# Tokeniser
text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
input_ids = text_input.input_ids.to(torch_device)

# Obtenir les enchâssements des tokens
token_embeddings = token_emb_layer(input_ids)

# Le nouvel enchâssement. Il s'agit maintenant d'un mélange d'enchâssement des tokens "puppy" et "skunk"
puppy_token_embedding = token_emb_layer(torch.tensor(6829, device=torch_device))
skunk_token_embedding = token_emb_layer(torch.tensor(42194, device=torch_device))
replacement_token_embedding = 0.5*puppy_token_embedding + 0.5*skunk_token_embedding

# Insérer ceci dans les enchâssements de token
token_embeddings[0, torch.where(input_ids[0]==6829)] = replacement_token_embedding.to(torch_device)

# Combiner avec le'enchâssement positionnel
input_embeddings = token_embeddings + position_embeddings

# Passage dans le transformer pour obtenir les enchâssements finaux
modified_output_embeddings = get_output_embeds(input_embeddings)

# Générer une image
generate_with_embs(modified_output_embeddings)

Inversion textuelle

Nous pouvons donc insérer un enchâssement de token modifié et l’utiliser pour générer une image. Nous avons utilisé l’enchâssement de token pour « chat » dans l’exemple ci-dessus, mais que se passerait-il si nous pouvions « apprendre » un nouvel enchâssement de token pour un concept spécifique ? C’est l’idée qui sous-tend l’« Inversion textuelle », dans laquelle quelques exemples d’images sont utilisés pour créer un nouvel enchâssement de token :

Diagramme tiré de l’article de blog sur l’inversion textuelle. Notez qu’il ne montre pas l’étape des enchâssements positionnels pour des raisons de simplicité.

Nous ne verrons pas comment cet entraînement fonctionne, mais nous pouvons essayer de charger l’un de ces nouveaux “concepts” à partir de la bibliothèque de concepts SD créée par la communauté et voir comment il s’intègre dans notre exemple ci-dessus. Nous utiliserons https://huggingface.co/sd-concepts-library/birb-style puisque c’est le premier que nous avons créé. Téléchargez le fichier learned_embeds.bin à partir de là et téléchargez-le à l’endroit où se trouve ce notebook avant d’exécuter la cellule suivante :

birb_embed = torch.load('learned_embeds.bin')
birb_embed.keys(), birb_embed['<birb-style>'].shape

(dict_keys(['<birb-style>']), torch.Size([768]))

Nous obtenons un dictionnaire avec une clé et l’enchâssement de token correspondant. Comme dans l’exemple précédent, remplaçons l’enchâssement de « puppy » par celui-ci et voyons ce qui se passe :

prompt = 'A mouse in the style of puppy'

# Tokeniser
text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
input_ids = text_input.input_ids.to(torch_device)

# Obtenir les enchâssements des tokens
token_embeddings = token_emb_layer(input_ids)

# Le nouvel enchâssement, notre mot d'ordre spécial
replacement_token_embedding = birb_embed['<birb-style>'].to(torch_device)

# Insérer ceci dans les enchâssements de token
token_embeddings[0, torch.where(input_ids[0]==6829)] = replacement_token_embedding.to(torch_device)

# Combiner avec le'enchâssement positionnel
input_embeddings = token_embeddings + position_embeddings

# Passage dans le transformer pour obtenir les enchâssements finaux
modified_output_embeddings = get_output_embeds(input_embeddings)

# Générer une image
generate_with_embs(modified_output_embeddings)

Le token a été remplacé par une expression qui représente un style particulier de peinture, mais il pourrait tout aussi bien représenter un objet ou une classe d’objets spécifique.

Encore une fois, il existe un beau notebook d’inférence d’Hugging Face pour faciliter l’utilisation des différents concepts, qui gère correctement l’utilisation des noms dans les prompts (“*A in the style of *") sans se préoccuper de toutes ces choses manuelles.

Mélanger les enchâssements

Outre le simple remplacement de l’enchâssement des tokens d’un seul mot, il existe d’autres astuces que nous pouvons essayer. Par exemple, que se passe-t-il si nous créons une « chimère » en calculant la moyenne des enchâssements de deux prompts différents ?

# Enchâsser deux prompts
text_input1 = tokenizer(["A mouse"], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
text_input2 = tokenizer(["A leopard"], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
    text_embeddings1 = text_encoder(text_input1.input_ids.to(torch_device))[0]
    text_embeddings2 = text_encoder(text_input2.input_ids.to(torch_device))[0]

# Les mixer ensemble
mix_factor = 0.35
mixed_embeddings = (text_embeddings1*mix_factor + \
                   text_embeddings2*(1-mix_factor))

# Generer
generate_with_embs(mixed_embeddings)

L’UNet et le CFG (Classifier Free Guidance)

Il est maintenant temps d’examiner le modèle de diffusion proprement dit. Il s’agit généralement d’un UNet qui prend en compte les latents bruyants (x) et prédit le bruit. Nous utilisons un modèle conditionnel qui prend également en compte le pas de temps (t) et notre enchâssement de texte (aka encoder_hidden_states) comme conditionnement. L’introduction de tous ces éléments dans le modèle se présente comme suit : noise_pred = unet(latents, t, encoder_hidden_states=text_embeddings)["sample"]

Nous pouvons l’essayer et voir à quoi ressemble le résultat :

# Preparation du planificateur
scheduler.set_timesteps(num_inference_steps)

# Quel est notre pas de temps ?
t = scheduler.timesteps[0]
sigma = scheduler.sigmas[0]

# Un latent bruyant
latents = torch.randn(
  (batch_size, unet.in_channels, height // 8, width // 8),
  generator=generator,
)
latents = latents.to(torch_device)
latents = latents * scheduler.init_noise_sigma

# L'enchâssement du texte
text_input = tokenizer(['A macaw'], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
    text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]

# Passage dans l'UNet pour prédire le bruit résiduel
with torch.no_grad():
    noise_pred = unet(latents, t, encoder_hidden_states=text_embeddings)["sample"]

latents.shape, noise_pred.shape # Nous obtenons des prédictions de la même forme que l'entrée

(torch.Size([1, 4, 64, 64]), torch.Size([1, 4, 64, 64]))

Étant donné un ensemble de latents bruyants, le modèle prédit la composante de bruit. Nous pouvons retirer ce bruit des latents bruyants pour voir à quoi ressemble l’image de sortie (latents_x0 = latents - sigma * noise_pred). Et nous pouvons ajouter la plus grande partie du bruit à cette sortie prédite pour obtenir l’entrée (légèrement moins bruitée, espérons-le) pour l’étape de diffusion suivante. Pour visualiser cela, générons une autre image, en sauvegardant à la fois la sortie prédite (x0) et l’étape suivante (xt-1) après chaque étape :

prompt = 'Oil painting of an otter in a top hat'
height = 512                      
width = 512                        
num_inference_steps = 50
guidance_scale = 8    
generator = torch.manual_seed(32)
batch_size = 1

# Créer un dossier pour stocker les résultats
!rm -rf steps/
!mkdir -p steps/

# Preparation du texte 
text_input = tokenizer([prompt], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
    text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
    [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
with torch.no_grad():
    uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0] 
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

# Preparation du planificateur
scheduler.set_timesteps(num_inference_steps)

# Preparation des latents
latents = torch.randn(
  (batch_size, unet.in_channels, height // 8, width // 8),
  generator=generator,
)
latents = latents.to(torch_device)
latents = latents * scheduler.init_noise_sigma

# Boucle
for i, t in tqdm(enumerate(scheduler.timesteps)):
    # étendre les latents si nous procédons à un guidage sans classifieur afin d'éviter de faire deux passages en avant
    latent_model_input = torch.cat([latents] * 2)
    sigma = scheduler.sigmas[i]
    latent_model_input = scheduler.scale_model_input(latent_model_input, t)

    # prédire le bruit résiduel
    with torch.no_grad():
        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]

    # effectuer le guidage
    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
    
    # Obtenir la valeur prédite x0 :
    # latents_x0 = latents - sigma * noise_pred # Calculer nous-mêmes
    latents_x0 = scheduler.step(noise_pred, t, latents).pred_original_sample # Utilisation du planificateur (Diffuseurs 0.4 et plus)

    # calculer l'échantillon bruité précédent x_t -> x_t-1
    latents = scheduler.step(noise_pred, t, latents).prev_sample

    # Vers des images PIL
    im_t0 = latents_to_pil(latents_x0)[0]
    im_next = latents_to_pil(latents)[0]

    # Combinez les deux images et enregistrez-les pour une visualisation ultérieure
    im = Image.new('RGB', (1024, 512))
    im.paste(im_next, (0, 0))
    im.paste(im_t0, (512, 0))
    im.save(f'steps/{i:04}.jpeg')

# Réaliser et diffuser la vidéo sur l'état d'avancement (modifier la largeur à 1024 pour une pleine résolution)
!ffmpeg -v 1 -y -f image2 -framerate 12 -i steps/%04d.jpeg -c:v libx264 -preset slow -qp 18 -pix_fmt yuv420p out.mp4
mp4 = open('out.mp4','rb').read()
data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
HTML("""
<video width=600 controls>
      <source src="%s" type="video/mp4">
</video>
""" % data_url)

La version de droite montre la « sortie finale » prédite (x0) à chaque étape, et c’est ce qui est généralement utilisé pour les vidéos de progression, etc. La version de gauche représente l’étape suivante. Nous trouvons intéressant de comparer les deux, en regardant les vidéos de progression, on pourrait penser que des changements radicaux se produisent, en particulier aux premiers stades, mais comme les changements apportés à chaque étape sont relativement faibles, le processus réel est beaucoup plus progressif.

CFG (Classifier Free Guidance)

Par défaut, le modèle ne fait pas souvent ce que nous lui demandons. Si nous voulons qu’il suive mieux le prompt, nous utilisons un hack appelé CFG. Il y a une bonne explication dans cette vidéo video d’AI Coffee Break with Letitia.

Dans le code, cela revient à faire :

noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

Cela fonctionne étonnamment bien :) Essayez de changer le guidance_scale dans le code ci-dessus et voyez comment cela affecte les résultats. Jusqu’où pouvez-vous aller avant que les résultats n’empirent ?

Échantillonnage

Il y a encore de la complexité cachée dans latents = scheduler.step(noise_pred, i, latents)["prev_sample"]. Comment l’échantillonneur passe-t-il exactement des latents bruyants actuels à une version légèrement moins bruyante ? Pourquoi ne pas utiliser le modèle en une seule étape ? Existe-t-il d’autres façons de voir les choses ?

Le modèle tente de prédire le bruit dans une image. Pour des valeurs de bruit faibles, nous supposons qu’il fait un assez bon travail. Pour des niveaux de bruit plus élevés, la tâche est ardue ! Ainsi, au lieu de produire une image parfaite, les résultats ont tendance à ressembler à un désordre flou. Voir le début de la vidéo citée à l’instant pour une illustration ! Les échantillonneurs utilisent donc les prédictions du modèle pour s’en rapprocher légèrement (en éliminant une partie du bruit), puis obtiennent une autre prédiction basée sur cette entrée marginalement moins mauvaise, en espérant que cela améliorera le résultat de manière itérative.

Les différents échantillonneurs procèdent de différentes manières. Vous pouvez essayer d’inspecter le code de l’échantillonneur LMS par défaut avec :

# ??scheduler.step

Guidage

Ok, dernière astuce ! Comment pouvons-nous ajouter un contrôle supplémentaire à ce processus de génération ?

À chaque étape, nous allons utiliser notre modèle comme précédemment pour prédire la composante bruit de $x$. Ensuite, nous allons l’utiliser pour produire une image de sortie prédite, et appliquer une fonction de perte à cette image.

Cette fonction peut être n’importe quoi, mais nous allons faire une démonstration avec un exemple très simple. Si nous voulons des images avec beaucoup de bleu, nous pouvons créer une fonction de perte qui donne une perte élevée si les pixels ont une faible composante bleue :

def blue_loss(images):
    # Quelle est la distance entre les valeurs du canal bleu et 0,9 ?
    error = torch.abs(images[:,2] - 0.9).mean() # [:,2] -> toutes les images dans le batch, seulement le canal bleu
    return error

Lors de chaque étape de mise à jour, nous trouvons le gradient de la perte par rapport aux latents bruyants actuels et nous les modifions dans la direction qui réduit cette perte tout en effectuant l’étape de mise à jour normale :

prompt = 'A campfire (oil on canvas)' #@param
height = 512                        # hauteur par défaut de Stable Diffusion
width = 512                         # largeur par défaut de Stable Diffusion
num_inference_steps = 50  #@param           # Nombre d'étapes de débruitage
guidance_scale = 8 #@param               # Échelle pour un guidage sans classifieur
generator = torch.manual_seed(32)   # Générateur de graines pour créer le bruit latent initial
batch_size = 1
blue_loss_scale = 200 #@param

# Preparation du texte 
text_input = tokenizer([prompt], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
    text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]

# Et l'entrée non conditionnelle comme précédemment :
max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
    [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
with torch.no_grad():
    uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0] 
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

# Preparation du planificateur
scheduler.set_timesteps(num_inference_steps)

# Preparation des latents
latents = torch.randn(
  (batch_size, unet.in_channels, height // 8, width // 8),
  generator=generator,
)
latents = latents.to(torch_device)
latents = latents * scheduler.init_noise_sigma

# Boucle
for i, t in tqdm(enumerate(scheduler.timesteps)):
    # étendre les latents si nous procédons à un guidage sans classifieur afin d'éviter de faire deux passages en avant
    latent_model_input = torch.cat([latents] * 2)
    sigma = scheduler.sigmas[i]
    latent_model_input = scheduler.scale_model_input(latent_model_input, t)

    # prédire le bruit résiduel
    with torch.no_grad():
        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]

    # réaliser le CFG
    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
    
    #### GUIDAGES SUPPLÉMENTAIRES ###
    if i%5 == 0:
        # Requires_grad sur les latents
        latents = latents.detach().requires_grad_()

        # Obtenir la valeur prédite x0 :
        # latents_x0 = latents - sigma * noise_pred
        latents_x0 = scheduler.step(noise_pred, t, latents).pred_original_sample

        # Décodage vers l'espace d'image
        denoised_images = vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5 # range (0, 1)

        # Calculer la perte
        loss = blue_loss(denoised_images) * blue_loss_scale
        
        # Imprimer occasionnellement
        if i%10==0:
            print(i, 'loss:', loss.item())

        # Obtenir le gradient
        cond_grad = torch.autograd.grad(loss, latents)[0]

        # Modifier les latents en fonction de ce gradient
        latents = latents.detach() - cond_grad * sigma**2
    
    # Etape avec le planificateur
    latents = scheduler.step(noise_pred, t, latents).prev_sample


latents_to_pil(latents)[0]

0 loss: 182.02133178710938 
10 loss: 43.55351257324219 
20 loss: 15.30621337890625 
30 loss: 9.746519088745117 
40 loss: 8.846868515014648

Ajustez l’échelle (blue_loss_scale). A faible valeur, l’image est principalement rouge et orange grâce au prompt. Avec des valeurs plus élevées, l’image est surtout bleutée ! Si l’échelle est trop élevée, l’image devient bleue.

Comme c’est lent, vous remarquerez que nous appliquons cette perte seulement une fois toutes les 5 itérations. Pour vos propres tests, vous pouvez envisager d’utiliser une échelle plus basse pour la perte et de l’appliquer à chaque itération à la place :)

NB : Nous devrions mettre latents requires_grad=True avant de faire la passe avant de l’UNet (en enlevant avec torch.no_grad()) si nous voulons des gradients précis. MAIS cela nécessite beaucoup de mémoire supplémentaire. Vous verrez les deux approches utilisées en fonction de l’implémentation que vous regardez.

Le guidage avec des modèles de classification peut vous donner des images d’une classe spécifique. L’utilisation d’un modèle comme CLIP permet de mieux faire correspondre un prompt à un texte. L’utilisation d’une perte de style permet d’ajouter un style particulier. Le guidage avec une sorte de perte de perception peut l’orienter vers l’aspect général d’une image cible. Et ainsi de suite.

Conclusion

Nous espérons que vous avez maintenant une idée un peu plus précise de ce qui se passe lorsque vous créez une image avec l’un de ces modèles, et de la façon dont vous pouvez modifier le processus de manière créative. En espérant que cela vous donne envie de faire quelque chose d’amusant :)

Ce notebook a été écrit par Jonathan Whitaker, adapté de Grokking Stable Diffusion qui était sa première tentative pour comprendre ces composants par lui-même. Si vous repérez des bugs ou si vous avez des questions, n’hésitez pas à le contacter via @johnowhitaker :) Bonne lecture !

4.1. Débruitage inverse des modèles de diffusion implicites (DDIM)

Dans ce notebook, nous allons explorer l’inversion, voir comment elle est liée à l’échantillonnage, et l’appliquer à la tâche d’édition d’images avec Stable Diffusion. Ce que vous allez apprendre :

• Comment fonctionne l’échantillonnage DDIM
• Échantillonneurs déterministes et stochastiques
• La théorie derrière l’inversion DDIM
• L’édition d’images avec l’inversion

Commençons !

Configuration

# !pip install -q transformers diffusers accelerate

import torch
import requests
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from io import BytesIO
from tqdm.auto import tqdm
from matplotlib import pyplot as plt
from torchvision import transforms as tfms
from diffusers import StableDiffusionPipeline, DDIMScheduler

# Une fonction utile pour plus tard
def load_image(url, size=None):
    response = requests.get(url,timeout=0.2)
    img = Image.open(BytesIO(response.content)).convert('RGB')
    if size is not None:
        img = img.resize(size)
    return img

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Chargement d’un pipeline existant

# Charger un pipeline
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)

# Mettre en place un planificateur DDIM
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)

# Échantillon d'une image pour s'assurer que tout fonctionne bien
prompt = 'Beautiful DSLR Photograph of a penguin on the beach, golden hour'
negative_prompt = 'blurry, ugly, stock photo'
im = pipe(prompt, negative_prompt=negative_prompt).images[0]
im.resize((256, 256)) # redimensionner pour une meilleure visualisation

Echantillonage DDIM

À un moment donné $t$, l’image bruitée $x_t$ est un mélange de l’image originale ($x_0$) et de bruit ($\epsilon$). Voici la formule pour $x_t$ tirée de l’article DDIM, à laquelle nous nous référerons dans cette section :

$\epsilon$ est un bruit gaussien de variance unitaire $\alpha_t$ (“alpha”) est la valeur qui est appelée de manière confuse $\bar{\alpha}$ (“alpha_bar”) dans le papier DDPM ( !!) et qui définit le planificateur de bruit. Dans 🤗 Diffusers, le planificateur alpha est calculé et les valeurs sont stockées dans le scheduler.alphas_cumprod. Je sais que c’est déroutant ! Traçons ces valeurs, et n’oubliez pas que pour le reste de ce notebook, nous utiliserons la notation de DDIM.

# Tracer 'alpha' (alpha_bar dans DDPM, alphas_cumprod dans Diffusers)
timesteps = pipe.scheduler.timesteps.cpu()
alphas = pipe.scheduler.alphas_cumprod[timesteps]
plt.plot(timesteps, alphas, label='alpha_t');
plt.legend()

Au départ (étape 0, côté gauche du graphique), nous commençons avec une image propre et sans bruit. $\alpha_t = 1$. Au fur et à mesure que nous passons à des pas de temps plus élevés, nous nous retrouvons avec presque tout le bruit et $\alpha_t$ chute vers 0.

Lors de l’échantillonnage, nous commençons avec du bruit pur au pas de temps $1000$ et nous nous rapprochons lentement du pas de temps $0$. Pour calculer le prochain $t$ de la trajectoire d’échantillonnage ($x_{t-1}$ puisque nous passons d’un $t$ élevé à un $t$ faible), nous prédisons le bruit ($\epsilon_\theta(x_t)$, qui est la sortie de notre modèle) et nous l’utilisons pour calculer l’image débruitée prédite $x_0$. Nous utilisons ensuite cette prédiction pour nous déplacer sur une petite distance dans la « direction pointant vers $x_t$ ». Enfin, nous pouvons ajouter du bruit supplémentaire à l’échelle de $\sigma_t$. Voici la section de l’article qui montre cette méthode en action :

Nous disposons donc d’une équation permettant de passer de $x_t$ à $x_{t-1}$, avec une quantité de bruit contrôlable. Dans notre cas présent, nous nous intéressons plus particulièrement au cas où nous n’ajoutons aucun bruit supplémentaire, ce qui nous donne un échantillonnage DDIM entièrement déterministe. Voyons ce que cela donne en code :

# Fonction d'échantillonnage (DDIM standard)
@torch.no_grad()
def sample(prompt, start_step=0, start_latents=None,
           guidance_scale=3.5, num_inference_steps=30,
           num_images_per_prompt=1, do_classifier_free_guidance=True,
           negative_prompt='', device=device):
  
    # Encoder le prompt
    text_embeddings = pipe._encode_prompt(
            prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt
    )

    # Nombre d'étapes d'inférence
    pipe.scheduler.set_timesteps(num_inference_steps, device=device)

    # Créer un point de départ aléatoire si nous n'en avons pas déjà un
    if start_latents is None:
        start_latents = torch.randn(1, 4, 64, 64, device=device)
        start_latents *= pipe.scheduler.init_noise_sigma

    latents = start_latents.clone()

    for i in tqdm(range(start_step, num_inference_steps)):
    
        t = pipe.scheduler.timesteps[i]

        # développer les latents si l'on procède à un guidage sans classifieur
        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
        latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)

        # prédire le bruit résiduel
        noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

        #réaliser un guidage
        if do_classifier_free_guidance:
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)


        # Normalement, nous devrions nous fier au planificateur pour gérer l'étape de mise à jour :
        # latents = pipe.scheduler.step(noise_pred, t, latents).prev_sample

        # Au lieu de cela, faisons-le nous-mêmes :
        prev_t = max(1, t.item() - (1000//num_inference_steps)) # t-1
        alpha_t = pipe.scheduler.alphas_cumprod[t.item()]
        alpha_t_prev = pipe.scheduler.alphas_cumprod[prev_t]
        predicted_x0 = (latents - (1-alpha_t).sqrt()*noise_pred) / alpha_t.sqrt()
        direction_pointing_to_xt = (1-alpha_t_prev).sqrt()*noise_pred
        latents = alpha_t_prev.sqrt()*predicted_x0 + direction_pointing_to_xt

    # Post-traitement
    images = pipe.decode_latents(latents)
    images = pipe.numpy_to_pil(images)

    return images

# Tester notre fonction d'échantillonnage en générant une image
sample('Watercolor painting of a beach sunset', negative_prompt=negative_prompt, num_inference_steps=50)[0].resize((256, 256))

Voyez si vous pouvez faire correspondre le code avec l’équation de l’article. Notez que $\sigma$=0 puisque nous ne nous intéressons qu’au cas où il n’y a pas de bruit supplémentaire, nous pouvons donc laisser de côté ces éléments de l’équation.

Inversion

L’objectif est d’inverser le processus d’échantillonnage. Nous voulons obtenir un latent bruité qui, s’il est utilisé comme point de départ de notre procédure d’échantillonnage habituelle, permet de générer l’image originale.

Ici, nous chargeons une image comme image initiale, mais vous pouvez également en générer une vous-même pour l’utiliser à la place.

# https://www.pexels.com/photo/a-beagle-on-green-grass-field-8306128/
input_image = load_image('https://images.pexels.com/photos/8306128/pexels-photo-8306128.jpeg', size=(512, 512))
input_image

Nous allons également utiliser un prompt pour effectuer l’inversion avec l’aide d’un classifieur libre, alors entrez une description de l’image :

input_image_prompt = "Photograph of a puppy on the grass"

Ensuite, nous devons transformer cette image PIL en un ensemble de latents que nous utiliserons comme point de départ de notre inversion :

# encoder avec le VAE
with torch.no_grad(): latent = pipe.vae.encode(tfms.functional.to_tensor(input_image).unsqueeze(0).to(device)*2-1)
l = 0.18215 * latent.latent_dist.sample()

Très bien, il est temps de passer à la partie amusante. Cette fonction ressemble à la fonction d’échantillonnage ci-dessus, mais nous nous déplaçons à travers les pas de temps dans la direction opposée, en commençant à $t=0$ et en nous déplaçant vers un bruit de plus en plus élevé. Et au lieu de mettre à jour nos latents pour qu’ils soient moins bruyants, nous estimons le bruit prédit et l’utilisons pour ANNULER une étape de mise à jour, en les déplaçant de $t$ à $t+1$.

## Inversion
@torch.no_grad()
def invert(start_latents, prompt, guidance_scale=3.5, num_inference_steps=80,
           num_images_per_prompt=1, do_classifier_free_guidance=True,
           negative_prompt='', device=device):
  
    # Encoder le prompt
    text_embeddings = pipe._encode_prompt(
            prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt
    )

    # les latents sont maintenant les latents de départ spécifiés
    latents = start_latents.clone()

    # Nous garderons une liste des latents inversés au fur et à mesure du processus
    intermediate_latents = []

    # Définir le nombre d'étapes de l'inférence
    pipe.scheduler.set_timesteps(num_inference_steps, device=device)

    # Pas de temps inversés <<<<<<<<<<<<<<<<<<<<
    timesteps = reversed(pipe.scheduler.timesteps)

    for i in tqdm(range(1, num_inference_steps), total=num_inference_steps-1):

        # Nous allons sauter l'itération finale
        if i >= num_inference_steps - 1: continue

        t = timesteps[i]

		# développer les latents si l'on fait de l'orientation sans classifieur
        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
        latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)

        # prédire le bruit résiduel
        noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

        # effectuer un guidage
        if do_classifier_free_guidance:
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

        current_t = max(0, t.item() - (1000//num_inference_steps))#t
        next_t = t # min(999, t.item() + (1000//num_inference_steps)) # t+1
        alpha_t = pipe.scheduler.alphas_cumprod[current_t]
        alpha_t_next = pipe.scheduler.alphas_cumprod[next_t]

        # Étape de mise à jour inversée (réorganisation de l'étape de mise à jour pour obtenir x(t) (nouveaux latents) en fonction de x(t-1) (latents actuels)
        latents = (latents - (1-alpha_t).sqrt()*noise_pred)*(alpha_t_next.sqrt()/alpha_t.sqrt()) + (1-alpha_t_next).sqrt()*noise_pred


        # Stockage
        intermediate_latents.append(latents)
            
    return torch.cat(intermediate_latents)

En l’exécutant sur la représentation latente de notre photo de chiot, nous obtenons un ensemble de tous les latents intermédiaires créés au cours du processus d’inversion :

inverted_latents = invert(l, input_image_prompt,num_inference_steps=50)
inverted_latents.shape

torch.Size([48, 4, 64, 64])

Nous pouvons visualiser l’ensemble final de latents - ceux-ci constitueront, nous l’espérons, le point de départ bruyant de nos nouvelles tentatives d’échantillonnage :

# Décoder les latents inversés finaux
with torch.no_grad():
  im = pipe.decode_latents(inverted_latents[-1].unsqueeze(0))
pipe.numpy_to_pil(im)[0]

Vous pouvez transmettre ces latents inversés au pipeline en utilisant la méthode__call__ normale :

pipe(input_image_prompt, latents=inverted_latents[-1][None], num_inference_steps=50, guidance_scale=3.5).images[0]

Mais c’est là que nous voyons notre premier problème : ce n’est pas tout à fait l’image avec laquelle nous avons commencé ! En effet, l’inversion DDIM repose sur une hypothèse critique selon laquelle la prédiction du bruit à l’instant $t$ et à l’instant $t+1$ sera la même, ce qui n’est pas vrai lorsque l’inversion ne porte que sur $50$ ou $100$ pas de temps. Nous pourrions utiliser davantage de pas de temps pour espérer obtenir une inversion plus précise, mais nous pouvons également tricher et commencer à partir de, disons, $20/50$ pas d’échantillonnage avec les latents intermédiaires correspondants que nous avons sauvegardés lors de l’inversion :

# La raison pour laquelle nous voulons pouvoir spécifier l'étape de démarrage
start_step=20
sample(input_image_prompt, start_latents=inverted_latents[-(start_step+1)][None], 
       start_step=start_step, num_inference_steps=50)[0]

Très proche de notre image d’entrée ! Pourquoi faisons-nous cela ? Eh bien, l’espoir est que si nous échantillonnons maintenant avec un nouveau prompt, nous obtiendrons une image qui correspond à l’original SAUF aux endroits pertinents pour le nouveau prompt. Par exemple, en remplaçant « puppy » par « cat », nous devrions voir un chat avec un dos et un arrière-plan presque identiques :

# Échantillonnage avec un nouveau prompt
start_step=10
new_prompt = input_image_prompt.replace('puppy', 'cat')
sample(new_prompt, start_latents=inverted_latents[-(start_step+1)][None], 
       start_step=start_step, num_inference_steps=50)[0]

Pourquoi ne pas utiliser img2img ?

Pourquoi s’embêter à inverser ? Ne peut-on pas simplement ajouter du bruit à l’image d’entrée et la débruiter avec le nouveau prompt ? Nous le pouvons, mais cela entraînera des changements beaucoup plus radicaux partout (si nous ajoutons beaucoup de bruit) ou des changements insuffisants partout (si nous ajoutons moins de bruit). Essayez vous-même :

start_step = 10
num_inference_steps=50
pipe.scheduler.set_timesteps(num_inference_steps)
noisy_l = pipe.scheduler.add_noise(l, torch.randn_like(l), pipe.scheduler.timesteps[start_step])
sample(new_prompt, start_latents=noisy_l, start_step=start_step, num_inference_steps=num_inference_steps)[0]

Notez la modification beaucoup plus importante de la pelouse et de l’arrière-plan.

Rassembler le tout

Rassemblons le code que nous avons écrit jusqu’à présent dans une fonction simple qui prend une image et deux prompts et effectue une modification en utilisant l’inversion :

def edit(input_image, input_image_prompt, edit_prompt, num_steps=100, start_step=30, guidance_scale=3.5):
    with torch.no_grad(): latent = pipe.vae.encode(tfms.functional.to_tensor(input_image).unsqueeze(0).to(device)*2-1)
    l = 0.18215 * latent.latent_dist.sample()
    inverted_latents = invert(l, input_image_prompt,num_inference_steps=num_steps)
    final_im = sample(edit_prompt, start_latents=inverted_latents[-(start_step+1)][None], 
                      start_step=start_step, num_inference_steps=num_steps, guidance_scale=guidance_scale)[0]
    return final_im

Et en action :

edit(input_image, 'A puppy on the grass', 'an old grey dog on the grass', num_steps=50, start_step=10)

edit(input_image, 'A puppy on the grass', 'A blue dog on the lawn', num_steps=50, start_step=12, guidance_scale=6)

✏️ À votre tour ! Essayez ceci sur d’autres images ! Explorez les différents paramètres.

Plus de pas = meilleure performance

Si vous avez des problèmes avec des inversions moins précises, vous pouvez essayer d’utiliser plus de pas (au prix d’un temps d’exécution plus long). Pour tester l’inversion, vous pouvez utiliser notre fonction d’édition avec le même prompt :

# Test d'inversion avec beaucoup plus d'étapes :
edit(input_image, 'A puppy on the grass', 'A puppy on the grass', num_steps=350, start_step=1)

C’est beaucoup mieux ! Et en essayant de l’éditer :

edit(input_image, 'A photograph of a puppy', 'A photograph of a grey cat', num_steps=150, start_step=30, guidance_scale=5.5)

# source: https://www.pexels.com/photo/girl-taking-photo-1493111/
face = load_image('https://images.pexels.com/photos/1493111/pexels-photo-1493111.jpeg', size=(512, 512))
face

edit(face, 'A photograph of a face', 'A photograph of a face with sunglasses', num_steps=250, start_step=30, guidance_scale=3.5)

edit(face, 'A photograph of a face', 'Acrylic palette knife painting of a face, colorful', num_steps=250, start_step=65, guidance_scale=5.5)

Et ensuite ?

Armé des connaissances de ce notebook, nous vous recommandons d’étudier Null-text Inversion qui s’appuie sur DDIM en optimisant le texte nul (prompt inconditionnel) lors de l’inversion pour des inversions plus précises et de meilleures éditions.

4.2. Diffusion pour l'audio

Dans ce notebook, nous allons jeter un bref coup d’œil à la génération d’audio avec des modèles de diffusion. Ce que vous allez apprendre :

• Comment l’audio est représenté dans un ordinateur
• Les méthodes de conversion entre les données audio brutes et les spectrogrammes
• Comment préparer un chargeur de données avec une fonction personnalisée pour convertir des tranches d’audio en spectrogrammes
• Finetuner un modèle de diffusion audio existant sur un genre de musique spécifique
• Télécharger votre pipeline personnalisé sur le Hub d’Hugging Face

Mise en garde : il s’agit principalement d’un objectif pédagogique - rien ne garantit que notre modèle sonnera bien 😉

Commençons !

Configuration et importations

# !pip install -q datasets diffusers torchaudio accelerate

import torch, random
import numpy as np
import torch.nn.functional as F
from tqdm.auto import tqdm
from IPython.display import Audio
from matplotlib import pyplot as plt
from diffusers import DiffusionPipeline
from torchaudio import transforms as AT
from torchvision import transforms as IT

Echantillonnage à partir d’un pipeline audio pré-entraîné

Commençons par suivre la documentation pour charger un modèle de diffusion audio préexistant :

# Chargement d'un pipeline de diffusion audio pré-entraîné
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = DiffusionPipeline.from_pretrained("teticio/audio-diffusion-instrumental-hiphop-256").to(device)

Comme pour les pipelines que nous avons utilisés dans les unités précédentes, nous pouvons créer des échantillons en appelant le pipeline comme suit :

# Échantillonner à partir du pipeline et afficher les résultats
output = pipe()
display(output.images[0])
display(Audio(output.audios[0], rate=pipe.mel.get_sample_rate()))

Ici, l’argument rate spécifie la fréquence d’échantillonnage de l’audio ; nous y reviendrons plus tard. Vous remarquerez également que le pipeline renvoie plusieurs choses. Que se passe-t-il ici ? Examinons de plus près les deux sorties.

La première est un tableau de données, représentant l’audio généré :

# Le tableau audio :
output.audios[0].shape

(1, 130560)

La seconde ressemble à une image en niveaux de gris :

# L'image de sortie (spectrogramme)
output.images[0].size

(256, 256)

Cela nous donne un aperçu du fonctionnement de ce pipeline. L’audio n’est pas directement généré par diffusion. Au lieu de cela, le pipeline a le même type d’UNet 2D que les pipelines de génération d’images inconditionnelles que nous avons vus dans l’unité 1, qui est utilisé pour générer le spectrogramme, qui est ensuite post-traité dans l’audio final.

Le pipeline possède un composant supplémentaire qui gère ces conversions, auquel nous pouvons accéder via pipe.mel :

pipe.mel

Mel {
  "_class_name": "Mel",
  "_diffusers_version": "0.12.0.dev0",
  "hop_length": 512,
  "n_fft": 2048,
  "n_iter": 32,
  "sample_rate": 22050,
  "top_db": 80,
  "x_res": 256,
  "y_res": 256
}

De l’audio à l’image et inversement

Une « forme d’onde » encode les échantillons audio bruts dans le temps. Il peut s’agir du signal électrique reçu d’un microphone, par exemple. Travailler avec cette représentation du « domaine temporel » peut s’avérer délicat, c’est pourquoi il est courant de la convertir sous une autre forme, communément appelée spectrogramme. Un spectrogramme montre l’intensité de différentes fréquences (axe y) en fonction du temps (axe x) :

# Calculer et afficher un spectrogramme pour notre échantillon audio généré en utilisant torchaudio
spec_transform = AT.Spectrogram(power=2)
spectrogram = spec_transform(torch.tensor(output.audios[0]))
print(spectrogram.min(), spectrogram.max())
log_spectrogram = spectrogram.log()
plt.imshow(log_spectrogram[0], cmap='gray');

tensor(0.) tensor(6.0842)

Le spectrogramme que nous venons de créer contient des valeurs comprises entre 0,0000000000001 et 1, la plupart d’entre elles étant proches de la limite inférieure de cette plage. Ce n’est pas l’idéal pour la visualisation ou la modélisation. En fait, nous avons dû prendre le logarithme de ces valeurs pour obtenir un tracé en niveaux de gris qui montre des détails. Pour cette raison, nous utilisons généralement un type spécial de spectrogramme appelé Mel spectrogramme, qui est conçu pour capturer les types d’informations qui sont importantes pour l’audition humaine en appliquant certaines transformations aux différentes composantes de fréquence du signal.

Quelques transformations audio de la documentation torchaudio

Heureusement pour nous, nous n’avons pas besoin de nous préoccuper de ces transformations, la fonctionnalité mel du pipeline s’occupe de ces détails pour nous. En l’utilisant, nous pouvons convertir une image de spectrogramme en audio comme suit :

a = pipe.mel.image_to_audio(output.images[0])
a.shape

(130560,)

Nous pouvons également convertir un tableau de données audio en images de spectrogramme en chargeant d’abord les données audio brutes, puis en appelant la fonction audio_slice_to_image(). Les clips plus longs sont automatiquement découpés en morceaux de la bonne longueur pour produire une image de spectrogramme de 256x256 :

pipe.mel.load_audio(raw_audio=a)
im = pipe.mel.audio_slice_to_image(0)
im

L’audio est représenté sous la forme d’un long tableau de nombres. Pour l’écouter nous avons besoin d’une autre information clé : la fréquence d’échantillonnage. Combien d’échantillons (valeurs individuelles) utilisons-nous pour représenter une seconde d’audio ?

Nous pouvons voir la fréquence d’échantillonnage utilisée lors de l’entraînement de ce pipeline avec :

sample_rate_pipeline = pipe.mel.get_sample_rate()
sample_rate_pipeline

22050

Si nous spécifions mal la fréquence d’échantillonnage, nous obtenons un son accéléré ou ralenti :

display(Audio(output.audios[0], rate=44100)) # Vitesse x2

Finetuning du pipeline

Maintenant que nous avons une compréhension approximative du fonctionnement du pipeline, nous allons le finetuner sur de nouvelles données audio !

Le jeu de données est une collection de clips audio de différents genres, que nous pouvons charger depuis le Hub de la manière suivante :

from datasets import load_dataset
dataset = load_dataset('lewtun/music_genres', split='train')
dataset

Dataset({
    features: ['audio', 'song_id', 'genre_id', 'genre'],
    num_rows: 19909
})

Vous pouvez utiliser le code ci-dessous pour voir les différents genres dans le jeu de données et combien d’échantillons sont contenus dans chacun d’eux :

for g in list(set(dataset['genre'])):
  print(g, sum(x==g for x in dataset['genre']))

Pop 945
Blues 58
Punk 2582
Old-Time / Historic 408
Experimental 1800
Folk 1214
Electronic 3071
Spoken 94
Classical 495
Country 142
Instrumental 1044
Chiptune / Glitch 1181
International 814
Ambient Electronic 796
Jazz 306
Soul-RnB 94
Hip-Hop 1757
Easy Listening 13
Rock 3095

Le jeu de données contient les données audio sous forme de tableaux :

audio_array = dataset[0]['audio']['array']
sample_rate_dataset = dataset[0]['audio']['sampling_rate']
print('Audio array shape:', audio_array.shape)
print('Sample rate:', sample_rate_dataset)
display(Audio(audio_array, rate=sample_rate_dataset))

Audio array shape: (1323119,)
Sample rate: 44100

Notez que la fréquence d’échantillonnage de cet audio est plus élevée. Si nous voulons utiliser le pipeline existant, nous devrons le « rééchantillonner » pour qu’il corresponde à la fréquence d’échantillonnage. Les clips sont également plus longs que ceux pour lesquels le pipeline est configuré. Heureusement, lorsque nous chargeons l’audio à l’aide de pipe.mel, il découpe automatiquement le clip en sections plus petites :

a = dataset[0]['audio']['array'] # Obtenir le tableau audio
pipe.mel.load_audio(raw_audio=a) # Le charger avec pipe.mel
pipe.mel.audio_slice_to_image(0) # Visualiser la première "tranche" sous forme de spectrogramme

Nous devons penser à ajuster le taux d’échantillonnage, car les données de ce jeu de données comportent deux fois plus d’échantillons par seconde :

sample_rate_dataset = dataset[0]['audio']['sampling_rate']
sample_rate_dataset

44100

Ici, nous utilisons les transformations de torchaudio (importées sous le nom AT) pour effectuer le rééchantillonnage, le pipeline mel pour transformer l’audio en image et les transformations de torchvision (importées sous le nom IT) pour transformer les images en tenseurs. Nous obtenons ainsi une fonction qui transforme un clip audio en un tenseur de spectrogramme que nous pouvons utiliser pour nous entraîner :

resampler = AT.Resample(sample_rate_dataset, sample_rate_pipeline, dtype=torch.float32)
to_t = IT.ToTensor()

def to_image(audio_array):
  audio_tensor = torch.tensor(audio_array).to(torch.float32)
  audio_tensor = resampler(audio_tensor)
  pipe.mel.load_audio(raw_audio=np.array(audio_tensor))
  num_slices = pipe.mel.get_number_of_slices()
  slice_idx = random.randint(0, num_slices-1) # Piocher une tranche aléatoire à chaque fois (à l'exception de la dernière tranche courte)
  im = pipe.mel.audio_slice_to_image(slice_idx) 
  return im

Nous utiliserons notre fonction to_image() dans le cadre d’une fonction collate personnalisée pour transformer notre jeu de données en un chargeur de données utilisable pour l’entraînement. La fonction collate définit la manière de transformer un batch d’exemples du jeu de données en un batch final de données prêtes à être entraînées. Dans ce cas, nous transformons chaque échantillon audio en une image de spectrogramme et nous empilons les tenseurs résultants :

def collate_fn(examples):
  # vers l'image -> vers le tenseur -> redimensionnement vers (-1, 1) -> empiler dans le batch
  audio_ims = [to_t(to_image(x['audio']['array']))*2-1 for x in examples]
  return torch.stack(audio_ims)

# Créer un jeu de données avec uniquement le genre de chansons 'Chiptune / Glitch'
batch_size=4 # 4 sur Colab, 12 sur A100
chosen_genre = 'Electronic' # <<< Essayer d'entraîner sur des genres différents <<<
indexes = [i for i, g in enumerate(dataset['genre']) if g == chosen_genre]
filtered_dataset = dataset.select(indexes)
dl = torch.utils.data.DataLoader(filtered_dataset.shuffle(), batch_size=batch_size, collate_fn=collate_fn, shuffle=True)
batch = next(iter(dl))
print(batch.shape)

torch.Size([4, 1, 256, 256])

NB : Vous devrez utiliser une taille de batch inférieure (par exemple 4) à moins que vous ne disposiez d’une grande quantité de vRAM GPU.

Boucle d’entraînement

Voici une boucle d’entraînement simple qui s’exécute à travers le chargeur de données pour quelques époques afin de finetuner le pipeline UNet. Vous pouvez également ignorer cette cellule et charger le pipeline avec le code de la cellule suivante.

epochs = 3
lr = 1e-4

pipe.unet.train()
pipe.scheduler.set_timesteps(1000)
optimizer = torch.optim.AdamW(pipe.unet.parameters(), lr=lr)

for epoch in range(epochs):
    for step, batch in tqdm(enumerate(dl), total=len(dl)):
        
        # Préparer les images d'entrée
        clean_images = batch.to(device)
        bs = clean_images.shape[0]

        # Échantillonner un pas de temps aléatoire pour chaque image
        timesteps = torch.randint(
            0, pipe.scheduler.num_train_timesteps, (bs,), device=clean_images.device
        ).long()

        # Ajouter du bruit aux images propres en fonction de l'ampleur du bruit à chaque étape
        noise = torch.randn(clean_images.shape).to(clean_images.device)
        noisy_images = pipe.scheduler.add_noise(clean_images, noise, timesteps)

        # Obtenir la prédiction du modèle
        noise_pred = pipe.unet(noisy_images, timesteps, return_dict=False)[0]

        # Calculer la perte
        loss = F.mse_loss(noise_pred, noise)
        loss.backward(loss)

        # Mise à jour des paramètres du modèle à l'aide de l'optimiseur
        optimizer.step()
        optimizer.zero_grad()

# OU : Charger la version entraînée précédemment
pipe = DiffusionPipeline.from_pretrained("johnowhitaker/Electronic_test").to(device)

output = pipe()
display(output.images[0])
display(Audio(output.audios[0], rate=22050))

# Créer un échantillon plus long en passant un tenseur de bruit de départ avec une forme différente
noise = torch.randn(1, 1, pipe.unet.sample_size[0],pipe.unet.sample_size[1]*4).to(device)
output = pipe(noise=noise)
display(output.images[0])
display(Audio(output.audios[0], rate=22050))

Ce ne sont pas les résultats les plus impressionnants mais c’est un début :) Essayez d’ajuster le taux d’apprentissage et le nombre d’époques, et partagez vos meilleurs résultats sur Discord pour que nous puissions nous améliorer ensemble !

Quelques éléments à prendre en compte

• Nous travaillons avec des images de spectrogrammes carrés de 256 pixels ce qui limite la taille de nos batchs. Pouvez-vous récupérer de l’audio de qualité suffisante à partir d’un spectrogramme de 128x128 ?
• Au lieu d’une augmentation aléatoire de l’image, nous choisissons à chaque fois des tranches différentes du clip audio, mais cela pourrait-il être amélioré avec différents types d’augmentation lorsque l’on s’entraîne pendant de nombreuses époques ?
• Comment pourrions-nous utiliser cette méthode pour générer des clips plus longs ? Peut-être pourriez-vous générer un clip de départ de 5 secondes, puis utiliser des idées inspirées de la complétion d’images (inpainting) pour continuer à générer des segments audio supplémentaires à partir du clip initial…
• Quel est l’équivalent d’une image à image dans ce contexte de diffusion de spectrogrammes ?

Pousser sur le Hub

Une fois que vous êtes satisfait de votre modèle, vous pouvez le sauvegarder et le transférer sur le Hub pour que d’autres personnes puissent en profiter :

from huggingface_hub import get_full_repo_name, HfApi, create_repo, ModelCard

# Choisir un nom pour le modèle
model_name = "audio-diffusion-electronic"
hub_model_id = get_full_repo_name(model_name)

# Sauvegarder le pipeline localement
pipe.save_pretrained(model_name)

# Inspecter le contenu du dossier
!ls {model_name}

mel  model_index.json  scheduler  unet

# Créer un dépôt
create_repo(hub_model_id)

# Télécharger les fichiers
api = HfApi()
api.upload_folder(
    folder_path=f"{model_name}/scheduler", path_in_repo="scheduler", repo_id=hub_model_id
)
api.upload_folder(
    folder_path=f"{model_name}/mel", path_in_repo="mel", repo_id=hub_model_id
)
api.upload_folder(folder_path=f"{model_name}/unet", path_in_repo="unet", repo_id=hub_model_id)
api.upload_file(
    path_or_fileobj=f"{model_name}/model_index.json",
    path_in_repo="model_index.json",
    repo_id=hub_model_id,
)

# Pousser une carte de modèle
content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-audio-generation
- diffusion-models-class
---

# Model Card for Unit 4 of the [Diffusion Models Class 🧨](https://github.com/huggingface/diffusion-models-class)

This model is a diffusion model for unconditional audio generation of music in the genre {chosen_genre}

## Usage

```python
from IPython.display import Audio
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("{hub_model_id}")
output = pipe()
display(output.images[0])
display(Audio(output.audios[0], rate=pipe.mel.get_sample_rate()))
```python
"""

card = ModelCard(content)
card.push_to_hub(hub_model_id)

Conclusion

Ce notebook vous a donné, nous l’espérons, un petit aperçu du potentiel de la génération audio. Consultez certaines des références liées à la vue d’ensemble de cette unité pour voir des méthodes plus fantaisistes et des échantillons stupéfiants qu’elles peuvent créer !

5.1. Hackathon DreamBooth 🏆

📣 Le hackathon est maintenant terminé et les gagnants ont été annoncés sur Discord. Vous êtes toujours invités à entraîner des modèles et à les soumettre au classement, mais nous n’offrirons plus de prix ou de certificats pour le moment.

Bienvenue au Hackathon DreamBooth !
C’est un événement communautaire où vous personnaliserez un modèle de Stable Diffusion en le finetunant sur une poignée de vos images. Pour ce faire, vous utiliserez une technique puissante appelée DreamBooth, qui permet d’implanter un sujet (par exemple votre animal de compagnie ou votre plat préféré) dans le domaine de sortie du modèle de sorte qu’il puisse être synthétisé avec un unique identifiant dans le prompt.

Cette compétition est composée de 5 thèmes, où chacun rassemblera des modèles appartenant aux catégories suivantes :

• Animal 🐨 : Utilisez ce thème pour générer des images de votre animal de compagnie ou de votre animal favori en train de se promener dans l’Acropole, de nager ou de voler dans l’espace.
• Science 🔬 : Utilisez ce thème pour générer des images synthétiques de galaxies, de protéines ou de tout autre domaine des sciences naturelles et médicales.
• Nourriture 🍔 : Utilisez ce thème pour finetuner un modèle sur votre plat ou cuisine préféré.
• Paysage 🏔 : Utilisez ce thème pour créer de magnifiques paysages de votre montagne, lac ou jardin préféré.
• Carte blanche 🔥 : Utilisez ce thème pour créer des modèles dans la catégorie de votre choix !

Nous attribuerons des prix aux 3 modèles les plus appréciés par thème, et vous êtes encouragés à soumettre autant de modèles que vous le souhaitez !

Pour commencer

Suivez les étapes ci-dessous pour participer à cet événement :

1. Rejoignez le serveur Discord d’Hugging Face et consultez le canal #dreambooth-hackathon pour vous tenir au courant du déroulement de l’événement.
2. Lancez et exécutez le notebook pour entraîner vos modèles en cliquant sur l’un des liens ci-dessous. Veillez à sélectionner l’exécution sur GPU pour vous assurer que vos modèles s’entraînent rapidement !

Notebook	Colab	Kaggle	Gradient	Studio Lab
Entraînement de DreamBooth

Note 👋 :
Le notebook DreamBooth utilise le checkpoint CompVis/stable-diffusion-v1-4 comme modèle de Diffusion Stable à finetuner. Cependant, vous êtes totalement libre d’utiliser n’importe quel checkpoint, il vous suffira d’ajuster le code pour charger les composants appropriés et le vérificateur de sécurité (s’il existe). Voici quelques modèles intéressants à finetuner :

• runwayml/stable-diffusion-v1-5
• prompthero/openjourney
• stabilityai/stable-diffusion-2
• hakurei/waifu-diffusion
• stabilityai/stable-diffusion-2-1
• nitrosocke/elden-ring-diffusion

Evaluation & Classement

Pour être dans la course aux prix, poussez un ou plusieurs modèles DreamBooth sur le Hub avec le tag dreambooth-hackathon dans la carte du modèle (exemple). Ce tag est créé automatiquement par le notebook de DreamBooth, mais vous devrez l’ajouter si vous exécutez vos propres scripts.

Les modèles sont évalués en fonction du nombre de j’aimes et vous pouvez suivre le classement de votre modèle sur le tableau de classement du hackathon :

• Classement DreamBooth

Chronologie

• 21 décembre 2022 : Date de démarrage
• 31 décembre 2022 : Date limite d’inscription à Colab Pro
• 22 janvier 2023 : Date limite de soumission (clôture du classement)
• 23-27 janvier 2023 : Annonce des gagnants de chaque thème

Tous les délais sont fixés à 23h59 UTC le jour correspondant, sauf indication contraire.

Prix

Nous attribuerons 3 prix par thème, les gagnants étant déterminés par les modèles ayant le plus de j’aimes sur le tableau de classement :

1ère place

• Abonnement Hugging Face Pro pour 1 an ou un bon d’achat de 100 $ pour le magasin de produits dérivés d’Hugging Face

2ème place

• Une copie du livre NLP with Transformers ou un bon d’achat de 50 $ pour le magasin de produits dérivés d’Hugging Face

3ème place

• Abonnement Hugging Face Pro pour 1 mois ou un bon d’achat de 15 $ pour le magasin de produits dérivés d’Hugging Face

Nous fournirons également un certificat d’achèvement à tous les participants qui soumettront au moins 1 modèle DreamBooth au hackathon 🔥.

Calculs

Google Colab sponsorisera cet événement en offrant des crédits Colab Pro à 100 participants (sélectionnés au hasard). Les crédits seront distribués en janvier 2023 et vous avez jusqu’au 31 décembre pour vous inscrire. Pour vous inscrire, veuillez remplir ce formulaire.

FAQ

Quelles sont les données autorisées pour le finetuning ?

Vous pouvez utiliser toutes les images qui vous appartiennent ou pour lesquelles une licence permissive vous autorise à le faire. Si vous souhaitez soumettre un modèle entraîné sur des visages (par exemple dans le cadre d’une carte blanche), nous vous recommandons d’utiliser vos propres traits. Idéalement, utilisez vos propres données lorsque vous le pouvez - nous aimerions voir vos animaux de compagnie ou vos paysages locaux préférés, et nous soupçonnons que les j’aimes et les prix iront à ceux qui font quelque chose de gentil et de personnel 😁.

D’autres techniques de finetuning comme l’inversion textuelle sont-elles autorisées ?

Absolument ! Bien que ce hackathon soit centré sur DreamBooth, vous êtes les bienvenus (et encouragés) à expérimenter d’autres techniques de finetuning. Cela signifie aussi que vous pouvez utiliser n’importe quel framework, code ou service qui vous aide à créer des modèles magnifiques pour que la communauté puisse en profiter 🥰.

5.2. Sprint Dreambooth en Keras

Cette paget résume toutes les informations pertinentes requises pour l’événement. 📋.

Introduction

Dreambooth est une technique de finetuning permettant d’enseigner de nouveaux concepts visuels à des modèles de diffusion conditionnés par le texte en utilisant seulement 3 à 5 images. Avec Dreambooth, vous pouvez générer des images drôles et réalistes de votre chien, de vous-même et de n’importe quel concept avec quelques images en utilisant Stable Diffusion. DreamBooth a été proposé dans DreamBooth : Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation par Ruiz et al. (2022)

Nous allons entraîner les modèles Dreambooth à l’aide de KerasCV et construire des démonstrations sur ces modèles.

Dates importantes

• Lancement de l’événement : 6 mars 2023
  Avec Nataniel Ruiz sur DreamBooth, François Chollet sur KerasCV et Apolinario Passos sur 🤗 Diffusers

• Début du sprint : 7 mars 2023
• Fin du sprint : 1er avril 2023
• Résultats : 7 avril 2023

Démarrer 🚀

Pour commencer, rejoignez-nous sur hf.co/join/discord et prenez le rôle #open-source, et rencontrez-nous sur le canal #keras-working-group.

Nous hébergerons nos démonstrations dans cette organisation sur Hugging Face Hub : keras-dreambooth, envoyez une demande ici si vous souhaitez soumettre une proposition 🙂

Nous allons :

1. Finetuner Stable Diffusion sur n’importe quel concept que nous voulons en utilisant Dreambooth,
2. Pousser le modèle vers le Hub d’Hugging Face,
3. Remplir la carte du modèle,
4. Construire une démo à partir du modèle.

Avertissement: Les modèles entraînés doivent être dans l’une des 4 catégories mentionnées dans la section Soumission. Veuillez y jeter un coup d’œil avant d’entraîner votre modèle.

Entraînement du modèle

Vous pouvez trouver le notebook ici (en anglais) et l’adapter à votre propre jeu de données.

Quelques inspirations pour le finetuning :

1. Lowpoly World : Ce modèle génère des mondes low poly 🤯🌍
2. Future Diffusion : Ce modèle génère des images dans des concepts de science-fiction futuristes 🤖
3. Fantasy sword : Ce modèle génère des épées pour des jeux à thème fantastique 🧙‍♂️

Si vous avez besoin de plus d’indications sur l’implémentation de Dreambooth avec Keras, vous pouvez consulter ce dépôt.

Dreambooth avec KerasCV

Pour l’instant, les options d’inférence et de déploiement de KerasCV sont limitées, et c’est là que la bibliothèque diffusers vient à la rescousse. Avec seulement quelques lignes de code, nous pouvons convertir un modèle KerasCV en un modèle diffusers et utiliser les pipelines diffusers pour effectuer l’inférence. Vous pouvez obtenir plus d’informations ici. Consultez aussi ce Space pour convertir votre modèle KerasCV en un modèle diffusers.

Les dépôts diffusers sur le Hub bénéficient d’une API d’inférence gratuite et de petits widgets dans la page du modèle où les utilisateurs peuvent jouer avec le modèle.

from diffusers import StableDiffusionPipeline

# point de contrôle de Stable Diffusion converti de KerasCV
model_ckpt = "sayakpaul/text-unet-dogs-kerascv_sd_diffusers_pipeline"
pipeline = StableDiffusionPipeline.from_pretrained(model_ckpt)
pipeline.to("cuda")

unique_id = "sks"
class_label = "dog"
prompt = f"A photo of {unique_id} {class_label} in a bucket"
image = pipeline(prompt, num_inference_steps=50).images[0]

Hébergement du modèle

À la fin du notebook vous verrez une section dédiée à l’hébergement, et une section séparée pour l’inférence. Nous utiliserons les fonctions de chargement et de poussée de modèles spécifiques à Keras de la bibliothèque huggingface_hub : push_to_hub_keras et from_pretrained_keras. Nous allons d’abord pousser le modèle en utilisant push_to_hub_keras. Une fois le modèle poussé, vous verrez que le modèle est hébergé avec une carte de modèle comme ci-dessous :

Pour mieux versionner les modèles, permettre la découvrabilité et la reproductibilité, nous allons remplir la carte de modèle. Cliquez sur *Edit model card*. Nous allons d’abord remplir la section Metadata de la fiche de modèle. Si votre modèle est entraîné avec un jeu de données du Hub, vous pouvez remplir la section des jeux de données avec le jeu de données. Nous allons remplir pipeline_tag avec text-to-image et choisir une licence pour notre modèle.

Ensuite, nous remplirons la partie markdown. Les hyperparamètres et le graphe sont automatiquement générés, nous pouvons donc écrire une courte explication pour la description, l’utilisation prévue et le jeu de données.

Vous pouvez trouver l’exemple de dépôt ci-dessous ici.

Démo

Nous allons utiliser Gradio pour construire nos démonstrations pour les modèles que nous avons entraînés. Avec la classe Interface, c’est simple :

from huggingface_hub import from_pretrained_keras
from keras_cv import models
import gradio as gr

sd_dreambooth_model = models.StableDiffusion(
    img_width=512, img_height=512
)
db_diffusion_model = from_pretrained_keras("merve/dreambooth_diffusion_model")
sd_dreambooth_model._diffusion_model = db_diffusion_model

# générer des images
def infer(prompt):
    generated_images = sd_dreambooth_model.text_to_image(
        prompt
    )
    return generated_images 
    
    
output = gr.Gallery(label="Outputs").style(grid=(2,2))

# la fonction de passage, le type d'entrée pour le prompt, la sortie pour les images multiples
gr.Interface(infer, inputs=["text"], outputs=[output]).launch()

Vous pouvez consulter le fichier app.py de l’application ci-dessous et le réutiliser pour votre modèle !

Dreambooth Submission - a Hugging Face Space par keras-dreambooth

Cette application génère des images d’un corgi 🐶

Hébergement de la démonstration sur Spaces

Une fois notre application terminée, nous pouvons créer un Space sur Hugging Face pour héberger notre application. Vous pouvez aller sur huggingface.co, cliquer sur votre profil en haut à droite et sélectionner “New Space”.

Nous pouvons nommer notre Space, choisir une licence et sélectionner « Gradio » comme Space SDK.

Après avoir créé le Space, vous pouvez utiliser soit les instructions ci-dessous pour cloner le dépôt localement, ajouter vos fichiers et pousser, OU, l’interface graphique pour créer les fichiers et écrire le code dans le navigateur.

Pour télécharger votre fichier, cliquez sur « Add File » et faites glisser/déposer votre fichier.

Enfin, nous devons créer un fichier appelé requirements.txt et ajouter les conditions du projet Dreambooth comme ci-dessous :

keras-cv
tensorflow
huggingface-hub

Et votre application devrait être opérationnelle !

Nous hébergerons nos modèles et nos Spaces sous cette organisation. Vous pouvez transporter vos modèles et Spaces dans l’onglet paramètres sous Rename or transfer this model et sélectionner keras-dreambooth dans le menu déroulant.

Si vous ne voyez pas keras-dreambooth dans la liste déroulante, il est probable que vous ne soyez pas membre de l’organisation. Utilisez ce lien pour demander à rejoindre l’organisation.

Soumission

Vous pouvez soumettre votre projet dans trois thèmes :

• Nature et animaux (nature)
• Univers de science-fiction/fantastique (sci-fi)
• Conscient (consentful) : Associez-vous à un artiste pour finetuner son style avec son consentement ! Assurez-vous d’inclure une référence au consentement explicite de l’artiste (par exemple un tweet) dans votre carte de modèle.
• Carte blanche (wild-card) : Si votre soumission appartient à une catégorie qui n’est pas mentionnée ci-dessus, n’hésitez pas à l’étiqueter avec wild-card afin que nous puissions l’évaluer en dehors de cette catégorie.

Ajoutez les catégories et leurs identifiants à votre carte modèle et ajoutez keras-dreambooth aux métadonnées dans la section des tags. Voici un exemple de carte de modèle. Toutes les soumissions seront compilées dans ce classement et classées en fonction du nombre de likes sur un espace donné afin de déterminer les gagnants.

Prix

Nous choisirons trois gagnants parmi les applications soumises, en fonction du nombre de likes accordés à un espace dans une catégorie donnée.

🛍️ Le premier remportera un bon d’achat de 100$ sur hf.co/shop ou un an d’abonnement à Hugging Face Pro

🛍️ La deuxième remportera un bon d’achat de 50$ sur hf.co/shop ou le livre Natural Language Processing with Transformers.

🛍️ Le troisième remportera un bon d’achat de 30$ sur hf.co/shop ou trois mois d’abonnement à Hugging Face Pro

5.3. Sprint ControlNet en JAX/Diffusers

Bienvenue au sprint communautaire en JAX/Diffusers ! L’objectif de ce sprint est de travailler sur des modèles de diffusion amusants et créatifs en utilisant JAX et Diffusers.

Lors de cet événement, nous créerons diverses applications avec des modèles de diffusion en JAX/Flax et Diffusers en utilisant des heures TPU gratuites généreusement fournies par Google Cloud.

Ce document présente toutes les informations importantes pour faire une soumission au sprint.

Organisation

Les participants peuvent proposer des idées pour un projet intéressant impliquant des modèles de diffusion. Des équipes de 3 à 5 personnes seront ensuite formées autour des projets les plus prometteurs et les plus intéressants. Assurez-vous de lire la section Communication pour savoir comment proposer des projets, commenter les idées de projet des autres participants et créer une équipe.

Pour aider chaque équipe à mener à bien son projet, nous organiserons des conférences données par des scientifiques et des ingénieurs de Google, de Hugging Face et de la communauté open source. Les conférences auront lieu le 17 avril. Assurez-vous d’assister aux conférences pour tirer le meilleur parti de votre participation ! Consultez la section Conférences pour avoir une vue d’ensemble des conférences, y compris l’orateur et l’heure de la conférence.

Chaque équipe bénéficiera ensuite d’un accès gratuit à une VM TPU v4-8 du 14 avril au 1er mai. De plus, nous fournirons un exemple d’entraînement en JAX/Flax et Diffusers pour entraîner un ControlNet afin de lancer votre projet. Nous fournirons également des exemples sur la façon de préparer les jeux de données. Pendant le sprint, nous nous assurerons de répondre à toutes les questions que vous pourriez avoir sur JAX/Flax et Diffusers et nous aiderons chaque équipe autant que possible !

Nous ne distribuerons pas de TPU pour les équipes composées d’un seul membre. Nous vous encourageons donc à rejoindre une équipe ou à trouver des coéquipiers pour votre idée.

À la fin du sprint, chaque soumission sera évaluée par un jury et les trois meilleures démonstrations recevront un prix. Consultez la section Comment soumettre une démo pour plus d’informations et de suggestions sur la manière de soumettre votre projet.

Note : Même si nous fournissons un exemple pour entraîner ControlNet, les participants peuvent proposer des idées qui n’impliquent pas du tout un ControlNet du moment qu’elles sont centrées sur les modèles de diffusion.

Dates importantes

• 29/03 : Annonce officielle de la semaine de la communauté.
• 31/03 : Commencez à former des groupes dans le canal #jax-diffusers-ideas sur Discord.
• 10/04 : Collecte des données.
• 13/04 & 14/04 & 17/04 : Conférences de lancement sur YouTube.
• 14/04 à 17/04 : Début de l’accès aux TPU.
• 01/05 : Fermeture de l’accès aux TPU.
• 08/05 : Annonce des 10 meilleurs projets et des prix.

Note : Nous accepterons les candidatures tout au long du sprint.

Communication

Toutes les communications importantes auront lieu sur notre serveur Discord. Rejoignez le serveur en utilisant ce lien. Après avoir rejoint le serveur, prenez le rôle Diffusers dans le canal #role-assignment et dirigez-vous vers le canal #jax-diffusers-ideas pour partager votre idée sous la forme d’un message de forum. Pour vous inscrire, remplissez le formulaire d’inscription et nous vous donnerons accès à deux canaux Discord supplémentaires pour les discussions et le support technique, ainsi qu’un accès aux TPU. Les annonces importantes de l’équipe Hugging Face, Flax/JAX et Google Cloud seront publiées sur le serveur.

Le serveur Discord sera le lieu central où les participants pourront publier leurs résultats, partager leurs expériences d’apprentissage, poser des questions et obtenir une assistance technique pour les divers obstacles qu’ils rencontrent.

Pour les problèmes liés à Flax/JAX, Diffusers, Datasets ou pour des questions spécifiques à votre projet, nous interagirons à travers les dépôts publics et les forums :

• Flax : Issues, Questions
• JAX : Issues, Questions
• 🤗 Diffusers : Issues, Questions
• 🤗 Dataset s: Issues, Questions
• Questions spécifiques aux projets : Elles peuvent être posées sur le canal #jax-diffusers-ideas sur Discord.
• Questions relatives au TPU : Canal #jax-diffusers-tpu-support sur Discord.
• Discussion générale : #jax-diffusers-sprint channel sur Discord. Vous aurez accès aux canaux #jax-diffusers-tpu-support et #jax-diffusers-sprint une fois que vous aurez été accepté pour participer au sprint.

Lorsque vous demandez de l’aide, nous vous encourageons à poster le lien vers le forum sur le serveur Discord, plutôt que de poster directement des issues ou des questions. De cette façon, nous nous assurons que tout le monde peut bénéficier de vos questions, même après la fin du sprint.

Note : Après le 10 avril, si vous vous êtes inscrit sur le formulaire Google, mais que vous n’êtes pas dans le canal Discord, veuillez laisser un message sur l’annonce officielle du forum et envoyer un ping à @mervenoyan, @sayakpaul, et @patrickvonplaten. Il se peut que nous prenions un jour pour traiter ces demandes.

Conférences

Nous avons invité d’éminents chercheurs et ingénieurs de Google, Hugging Face, et de la communauté open-source qui travaillent dans le domaine de l’IA générative. Nous mettrons à jour cette section avec des liens vers les conférences, alors gardez un œil ici ou sur Discord dans le canal diffusion models core-announcements et programmez vos rappels !

13 avril 2023

Intervenant	Sujet	Horaire	Video
Emiel Hoogeboom, Google Brain	Pixel-Space Diffusion models for High Resolution Images	4.00pm-4.40pm CEST / 7.00am-7.40am PST
Apolinário Passos, Hugging Face	Introduction to Diffusers library	4.40pm-5.20pm CEST / 7.40am-08.20am PST
Ting Chen, Google Brain	Diffusion++: discrete data and high-dimensional generation	5.45pm-6.25pm CEST / 08.45am-09.25am PST

14 avril 2023

Intervenant	Sujet	Horaire	Video
Tim Salimans, Google Brain	Efficient image and video generation with distilled diffusion models	4.00pm-4.40pm CEST / 7.00am-7.40am PST
Suraj Patil, Hugging Face	Masked Generative Models: MaskGIT/Muse	4.40pm-5.20pm CEST / 7.40am-08.20am PST
Sabrina Mielke, John Hopkins University	From stateful code to purified JAX: how to build your neural net framework	5.20pm-6.00pm CEST / 08.20am-09.00am PST

17 avril 2023

Intervenant	Sujet	Horaire	Video
Andreas Steiner, Google Brain	JAX & ControlNet	4.00pm-4.40pm CEST / 7.00am-7.40am PST
Boris Dayma, craiyon	DALL-E Mini	4.40pm-5.20pm CEST / 7.40am-08.20am PST
Margaret Mitchell, Hugging Face	Ethics of Text-to-Image	5.20pm-6.00pm CEST / 08.20am-09.00am PST

Données et prétraitement

Dans cette section, nous verrons comment construire votre propre jeu de données pour entraîner ControlNet.

Préparer un grand jeu de données local

Monter un disque

Si vous avez besoin d’espace supplémentaire, vous pouvez suivre ce guide pour créer un disque persistant, l’attacher à votre VM TPU et créer un répertoire pour monter le disque. Vous pouvez ensuite utiliser ce répertoire pour stocker votre jeu de données.

Par ailleurs, la VM TPU attribuée à votre équipe dispose d’un disque de stockage persistant de 3 To. Pour apprendre à l’utiliser, consultez ce guide.

Prétraitement des données

Nous montrons ici comment préparer un grand jeu de données pour entraîner un modèle ControlNet avec filtre de Canny. Plus précisément, nous fournissons un exemple de script qui :

• Sélectionne 1 million de paires image-texte à partir d’un jeu de données existant COYO-700M.
• Télécharge chaque image et utilise le filtre de Canny pour générer l’image de conditionnement.
• Crée un métafichier qui relie toutes les images et les images traitées à leurs légendes.

Utilisez la commande suivante pour exécuter le script de prétraitement des données de l’exemple. Si vous avez monté un disque sur votre TPU, vous devez placer vos fichiers train_data_dir et cache_dir sur le disque monté.

python3 coyo_1m_dataset_preprocess.py \
 --train_data_dir="/mnt/disks/persist/data" \
 --cache_dir="/mnt/disks/persist" \
 --max_train_samples=1000000 \
 --num_proc=16

Une fois le script exécuté, vous trouverez un dossier de données dans le répertoire train_data_dir spécifié avec la structure de dossier ci-dessous :

data
├── images
│   ├── image_1.png
│   ├── .......
│   └── image_1000000.jpeg
├── processed_images
│   ├── image_1.png
│   ├── .......
│   └── image_1000000.jpeg
└── meta.jsonl

Charger un jeu de données

Pour charger un jeu de données à partir du dossier de données que vous venez de créer, vous devez ajouter un script de chargement de jeu de données à votre dossier de données. Le script de chargement de données doit porter le même nom que le dossier. Par exemple, si votre dossier de données est data, vous devez ajouter un script de chargement de données nommé data.py. Nous fournissons un exemple de script de chargement de données que vous pouvez utiliser. Tout ce que vous avez à faire est de mettre à jour le DATA_DIR avec le chemin correct de votre dossier de données. Pour plus de détails sur l’écriture d’un script de chargement de données, reportez-vous à la documentation.

Une fois que le script de chargement de données est ajouté à votre dossier de données, vous pouvez le charger avec :

dataset = load_dataset("/mnt/disks/persist/data", cache_dir="/mnt/disks/persist" )

Notez que vous pouvez utiliser --train_data_dir pour passer le répertoire de votre dossier de données au script d’entraînement et générer votre jeu de données automatiquement pendant l’entraînement.

Pour les grands jeux de données, nous recommandons de générer le jeu de données une seule fois et de le sauvegarder sur le disque à l’aide de la commande

dataset.save_to_disk("/mnt/disks/persist/dataset")

Vous pouvez ensuite réutiliser le jeu de données sauvegardé pour votre entraînement en passant --load_from_disk.

Voici un exemple d’exécution d’un script d’entraînement qui chargera le jeu de données depuis le disque.

export MODEL_DIR="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="/mnt/disks/persist/canny_model"
export DATASET_DIR="/mnt/disks/persist/dataset"
export DISK_DIR="/mnt/disks/persist"

python3 train_controlnet_flax.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --train_data_dir=$DATASET_DIR \
 --load_from_disk \
 --cache_dir=$DISK_DIR \
 --resolution=512 \
 --learning_rate=1e-5 \
 --train_batch_size=2 \
 --revision="non-ema" \
 --from_pt \
 --max_train_steps=500000 \
 --checkpointing_steps=10000 \
 --dataloader_num_workers=16 

Préparer un jeu de données avec MediaPipe et Hugging Face

Nous fournissons un notebook ( ) qui vous montre comment préparer un jeu de données pour entraîner ControlNet en utilisant MediaPipe et Hugging Face. Plus précisément, dans le notebook, nous montrons :

• Comment tirer parti des solutions MediaPipe pour extraire les articulations du corps de la pose à partir des images d’entrée.
• Prédire les légendes en utilisant BLIP-2 à partir des images d’entrée en utilisant 🤗 Transformers.
• Construire et pousser le jeu de données final vers le Hugging Face Hub en utilisant 🤗 Datasets.

Vous pouvez vous référer au notebook pour créer vos propres jeux de données en utilisant d’autres solutions MediaPipe. Ci-dessous, nous listons toutes les solutions pertinentes :

• Détection des points de repère de pose
• Détection de repères de visage
• Segmentation de selfie

Entraîner ControlNet

C’est peut-être la partie la plus amusante et la plus intéressante de ce document, car nous vous montrons ici comment entraîner un modèle ControlNet personnalisé.

Note : Pour ce sprint, vous n’êtes PAS limité à entraîner des ControlNets. Nous fournissons ce script d’entraînement comme référence pour vous permettre de démarrer.

Pour un entraînement plus rapide sur les TPU et les GPU, vous pouvez tirer parti de l’exemple d’entraînement Flax. Suivez les instructions ci-dessus pour obtenir le modèle et le jeu de données avant d’exécuter le script.

Mise en place de la VM TPU

Avant de continuer avec le reste de cette section, vous devez vous assurer que l’adresse email que vous utilisez a été ajoutée au projet hf-flax sur Google Cloud Platform. Si ce n’est pas le cas, merci de nous le faire savoir sur le serveur Discord (vous pouvez taguer @sayakpaul, @merve, et @patrickvonplaten).

Dans ce qui suit, nous allons décrire comment le faire en utilisant une console standard, mais vous devriez également être en mesure de vous connecter à la VM TPU via des IDE, comme Visual Studio Code, etc.

1. Vous devez installer le Google Cloud SDK. Veuillez suivre les instructions sur https://cloud.google.com/sdk.

2. Une fois le Google Cloud SDK installé, vous devez configurer votre compte en exécutant la commande suivante. Assurez-vous que correspond à l'adresse gmail que vous avez utilisée pour vous inscrire à cet événement.

    gcloud config set account <your-email-adress>
   

3. Assurons-nous également que le bon projet est défini au cas où votre email serait utilisé pour plusieurs projets gcloud :

    gcloud config set project hf-flax
   

4. Ensuite, vous devez vous authentifier. Vous pouvez le faire en exécutant la commande

    gcloud auth login
   

   Vous devriez obtenir un lien vers un site web où vous pouvez authentifier votre compte gmail.

5. Enfin, vous pouvez établir un tunnel SSH dans la VM TPU ! Veuillez exécuter la commande suivante en réglant la “–zone” sur us-central2-b et sur le nom de la TPU qui vous a été envoyé par email par l’équipe de Hugging Face.

    gcloud alpha compute tpus tpu-vm ssh <tpu-name> --zone <zone> --project hf-flax
   

Cela devrait établir un tunnel SSH dans la VM TPU !

Note : Vous n’êtes PAS supposé avoir accès à la console Google Cloud. Aussi, il se peut que vous ne receviez pas de lien d’invitation pour rejoindre le projet hf-flax. Mais vous devriez tout de même pouvoir accéder à la VM TPU en suivant les étapes ci-dessus .

Note : Les VM TPU sont déjà attachées à des disques de stockage persistants (de 3 TB). Cela sera utile au cas où votre équipe souhaiterait entraîner localement un jeu de données volumineux. Le nom du disque de stockage devrait également figurer dans l’e-mail que vous avez reçu. Suivez cette section pour plus de détails.

Installation de JAX

Commençons par créer un environnement virtuel Python :

python3 -m venv <your-venv-name>

Nous pouvons activer l’environnement en lançant :

source ~/<your-venv-name>/bin/activate

Installez ensuite Diffusers et les dépendances d’entraînement de la bibliothèque :

pip install git+https://github.com/huggingface/diffusers.git

Ensuite, clonez ce dépôt et installez JAX, Flax et les autres dépendances :

git clone https://github.com/huggingface/community-events
cd community-events/jax-controlnet-sprint/training_scripts
pip install -U -r requirements_flax.txt

Pour vérifier que JAX a été correctement installé, vous pouvez exécuter la commande suivante :

import jax
jax.device_count()

Cela devrait afficher le nombre de cœurs de la TPU, qui devrait être de 4 sur une VM TPUv4-8. Si Python n’est pas capable de détecter le périphérique TPU, veuillez consulter la section des erreurs possibles plus bas pour des solutions.

Si vous souhaitez utiliser le logging Weights and Biases, vous devez également installer wandb maintenant :

pip install wandb

Note : Weights & Biases est gratuit pour les étudiants, les éducateurs et les chercheurs universitaires. Tous les participants à notre événement sont qualifiés pour obtenir un compte d’équipe académique Weights & Biases. Pour créer votre équipe, vous pouvez visiter le site https://wandb.ai/create-team et choisir le type d’équipe “Academic”. Pour plus d’informations sur la création et la gestion d’une équipe Weights & Biases, vous pouvez consulter le site https://docs.wandb.ai/guides/app/features/teams.

Exécution du script d’entraînement

Maintenant, téléchargeons deux images de conditionnement que nous utiliserons pour lancer la validation pendant l’entraînement afin de suivre nos progrès

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

Nous vous encourageons à stocker ou à partager votre modèle avec la communauté. Pour utiliser le Hub, veuillez vous connecter à votre compte Hugging Face, ou (en créer un si vous n’en avez pas déjà un) :

huggingface-cli login

Assurez-vous que les variables d’environnement MODEL_DIR, OUTPUT_DIR et HUB_MODEL_ID sont définies. Les variables OUTPUT_DIR et HUB_MODEL_ID spécifient où sauvegarder le modèle sur le Hub :

export MODEL_DIR="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="runs/fill-circle-{timestamp}"
export HUB_MODEL_ID="controlnet-fill-circle"

Et enfin, démarrez l’entraînement (assurez-vous d’être dans le répertoire jax-controlnet-sprint/training_scripts) !

python3 train_controlnet_flax.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" \
 --validation_steps=1000 \
 --train_batch_size=2 \
 --revision="non-ema" \
 --from_pt \
 --report_to="wandb" \
 --tracker_project_name=$HUB_MODEL_ID \
 --num_train_epochs=11 \
 --push_to_hub \
 --hub_model_id=$HUB_MODEL_ID

Notez que l’argument --from_pt convertira votre point de contrôle pytorch en flax. Cependant, il ne fonctionnera qu’avec les points de contrôle au format diffusers. Si votre MODEL_DIR ne contient pas de points de contrôle au format diffusers, vous ne pouvez pas utiliser l’argument --from_pt. Vous pouvez convertir vos points de contrôle ckpt ou safetensors au format diffusers en utilisant ce script.

Puisque nous avons passé l’argument --push_to_hub, il va automatiquement créer un repo de modèle sous votre compte Hugging Face basé sur $HUB_MODEL_ID. À la fin de l’entraînement, le point de contrôle final sera automatiquement stocké sur le Hub. Vous pouvez trouver un exemple de modèle ici.

Notre script d’entraînement fournit également un support limité pour le streaming de grands jeux de données à partir du Hub. Afin d’activer le streaming, il faut également définir --max_train_samples. Voici un exemple de commande (tiré de cet article de blog) :

export MODEL_DIR="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="runs/uncanny-faces-{timestamp}"
export HUB_MODEL_ID="controlnet-uncanny-faces"

python3 train_controlnet_flax.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=multimodalart/facesyntheticsspigacaptioned \
 --streaming \
 --conditioning_image_column=spiga_seg \
 --image_column=image \
 --caption_column=image_caption \
 --resolution=512 \
 --max_train_samples 100000 \
 --learning_rate=1e-5 \
 --train_batch_size=1 \
 --revision="flax" \
 --report_to="wandb" \
 --tracker_project_name=$HUB_MODEL_ID

Notez cependant que les performances des TPUs peuvent être limitées car le streaming avec datasets n’est pas optimisé pour les images. Pour assurer un débit maximal, nous vous encourageons à explorer les options suivantes :

• Webdataset
• TorchData
• TensorFlow Datasets

Lorsque vous travaillez avec un jeu de données plus important, vous pouvez avoir besoin d’exécuter le processus d’entraînement pendant une longue période et il est utile d’enregistrer des points de contrôle réguliers au cours du processus. Vous pouvez utiliser l’argument suivant pour activer les points de contrôle intermédiaires :

 --checkpointing_steps=500

Cela permet d’enregistrer le modèle entraîné dans des sous-dossiers du dossier output_dir. Le nom des sous-dossiers correspond au nombre d’étapes effectuées jusqu’à présent ; par exemple : un point de contrôle sauvegardé après 500 étapes d’entraînement serait sauvegardé dans un sous-dossier nommé 500

Vous pouvez alors commencer votre entraînement à partir de ce point de contrôle sauvegardé avec

 --controlnet_model_name_or_path="./control_out/500" 

Nous soutenons l’entraînement avec la stratégie de pondération Min-SNR proposée dans Efficient Diffusion Training via Min-SNR Weighting Strategy qui permet d’obtenir une convergence plus rapide en rééquilibrant la perte. Pour l’utiliser, il faut définir l’argument --snr_gamma. La valeur recommandée est 5.0.

Nous supportons également l’accumulation de gradient, technique qui vous permet d’utiliser une taille de batch plus grande que celle que votre machine serait normalement capable de mettre en mémoire. Vous pouvez utiliser l’argument gradient_accumulation_steps pour définir les étapes d’accumulation du gradient. L’auteur de ControlNet recommande d’utiliser l’accumulation de gradient pour obtenir une meilleure convergence. Pour en savoir plus voir ici.

Vous pouvez profiler votre code avec :

 --profile_steps==5

Reportez-vous à la documentation JAX sur le profilage. Pour inspecter la trace de profil, vous devez installer et démarrer Tensorboard avec le plugin de profil :

pip install tensorflow tensorboard-plugin-profile
tensorboard --logdir runs/fill-circle-100steps-20230411_165612/

Le profil peut alors être inspecté à l’adresse http://localhost:6006/#profile.

Parfois vous obtiendrez des conflits de version (messages d’erreur comme Duplicate plugins for name projector), ce qui signifie que vous devez désinstaller et réinstaller toutes les versions de Tensorflow/Tensorboard (par exemple avec pip uninstall tensorflow tf-nightly tensorboard tb-nightly tensorboard-plugin-profile && pip install tf-nightly tbp-nightly tensorboard-plugin-profile).

Notez que la fonctionnalité de débogage du plugin Tensorboard profile est toujours en cours de développement. Toutes les vues ne sont pas entièrement fonctionnelles, et par exemple le trace_viewer coupe les événements après 1M (ce qui peut résulter en la perte de toutes vos traces de périphériques si par exemple vous profilez l’étape de compilation par accident).

Dépannage de votre VM TPU

TRES IMPORTANT : Un seul processus peut accéder aux cœurs de la TPU à la fois. Cela signifie que si plusieurs membres de l’équipe essaient de se connecter aux cœurs de la TPU, vous obtiendrez des erreurs telles que :

libtpu.so already in used by another process. Not attempting to load libtpu.so in this process.

Nous recommandons à chaque membre de l’équipe de créer son propre environnement virtuel, mais une seule personne devrait exécuter les processus d’entraînement lourds. De plus, veuillez vous relayer lors de l’installation de la TPUv4-8 afin que tout le monde puisse vérifier que JAX est correctement installé.

Si les membres de votre équipe n’utilisent pas actuellement la TPU mais que vous obtenez toujours ce message d’erreur. Vous devez tuer le processus qui utilise la TPU avec :

kill -9 PID

vous devrez remplacer le terme « PID » par le PID du processus qui utilise TPU. Dans la plupart des cas, cette information est incluse dans le message d’erreur. Par exemple, si vous obtenez

The TPU is already in use by a process with pid 1378725. Not attempting to load libtpu.so in this process.

vous pouvez faire

kill -9 1378725

Vous pouvez également utiliser la commande suivante pour trouver les processus utilisant chacune des puces TPU (par exemple, /dev/accel0 est l’une des puces TPU)

sudo lsof -w /dev/accel0

Pour tuer tous les processus à l’aide de /dev/accel0, il faut

sudo lsof -t /dev/accel0 | xargs kill -9

Si Python n’est pas capable de détecter votre périphérique TPU (i.e. quand vous faites jax.device_count() et qu’il sort 0), cela peut être dû au fait que vous n’avez pas les droits d’accès aux logs tpu, ou que vous avez un fichier tpu lock qui traîne. Exécutez les commandes suivantes pour résoudre le problème

sudo rm -f /tmp/libtpu_lockfile

sudo chmod o+w /tmp/tpu_logs/

Comment faire une soumission

Pour faire une soumission complète, vous devez avoir les éléments suivants sur le Hub d’Hugging Face :

• Un dépôt de modèle avec les poids du modèle et la carte du modèle,
• (Facultatif) Un dépôt de jeu de données avec une carte de jeu de données,
• Un Space qui permet aux autres d’interagir avec votre modèle.

Pousser les poids du modèle et la carte du modèle vers le Hub

Si vous utilisez le script d’entraînement (train_controlnet_flax.py) fourni dans ce répertoire

L’activation de l’argument push_to_hub dans les arguments d’entraînement va :

• Créer un dépôt de modèles localement et à distance sur le Hub,
• Créer une carte de modèle et l’écrire dans le dépôt de modèles local,
• Sauvegarder votre modèle dans le référentiel de modèles local,
• Pousser le dépôt local vers le Hub.

Votre carte de modèle générée automatiquement ressemblera à ceci : .

Vous pouvez modifier la carte de modèle pour qu’elle soit plus informative. Les cartes de modèle qui sont plus informatives que les autres auront plus de poids lors de l’évaluation.

Si vous avez entraîné un modèle personnalisé et que vous n’avez pas utilisé le script

Vous devez vous authentifier avec huggingface-cli login comme indiqué ci-dessus. Si vous utilisez une des classes de modèles disponibles dans diffusers, sauvegardez votre modèle avec la méthode save_pretrained de votre modèle.

model.save_pretrained("path_to_your_model_repository")

Après avoir sauvegardé votre modèle dans un dossier, vous pouvez simplement utiliser le script ci-dessous pour pousser votre modèle vers le Hub :

from huggingface_hub import create_repo, upload_folder

create_repo("username/my-awesome-model")
upload_folder(
    folder_path="path_to_your_model_repository",
    repo_id="username/my-awesome-model"
)

Ceci poussera votre modèle vers Hub. Après avoir poussé cela, vous devez créer la carte de modèle vous-même. Vous pouvez utiliser l’interface graphique pour l’éditer.

Chaque carte de modèle se compose de deux sections, les métadonnées et le texte libre. Vous pouvez éditer les métadonnées à partir des sections dans l’interface graphique. Si vous avez sauvegardé votre modèle en utilisant save_pretrained, vous n’avez pas besoin de fournir pipeline_tag et library_name. Sinon, fournissez pipeline_tag, library_name et le jeu de données s’il existe sur Hugging Face Hub. En plus de cela, vous devez ajouter jax-diffusers-event à la section tags.

---
license: apache-2.0
library_name: diffusers
tags:
- jax-diffusers-event
datasets:
- red_caps
pipeline_tag: text-to-image
---

Créer notre Space

Rédiger notre application

Nous utiliserons Gradio pour créer nos applications. Gradio possède deux API principales : Interface et Blocks. Interface est une API de haut niveau qui vous permet de créer une interface avec quelques lignes de code, et Blocks est une API de plus bas niveau qui vous donne plus de flexibilité sur les interfaces que vous pouvez construire. Le code doit être inclus dans un fichier appelé app.py.

Essayons de créer une application ControlNet comme exemple. L’API Interface fonctionne simplement comme suit :

import gradio as gr

# La fonction d'inférence prend en compte le prompt, le prompt négatif et l'image
def infer(prompt, negative_prompt, image):
    # implémentez votre fonction d'inférence ici
    return output_image

# vous devez passer les entrées et les sorties en fonction de la fonction d'inférence
gr.Interface(fn = infer, inputs = ["text", "text", "image"], outputs = "image").launch()

Vous pouvez personnaliser votre interface en passant title, description et examples à la fonction Interface.

title = "ControlNet on Canny Filter"
description = "This is a demo on ControlNet based on canny filter."
# vous devez passer vos exemples en fonction de vos entrées
# chaque liste intérieure est un exemple, chaque élément de la liste correspondant à un composant des `inputs`.
examples = [["a cat with cake texture", "low quality", "cat_image.png"]]
gr.Interface(fn = infer, inputs = ["text", "text", "image"], outputs = "image",
            title = title, description = description, examples = examples, theme='gradio/soft').launch()

Votre interface ressemblera à ceci :

Avec les blocs, vous pouvez ajouter des marques, des onglets, des composants sous les colonnes et les lignes, etc. Supposons que nous ayons deux ControlNets et que nous voulions les inclure dans un Space. Nous les placerons sous différents onglets dans une démo comme ci-dessous :

import gradio as gr

def infer_segmentation(prompt, negative_prompt, image):
    # votre fonction d'inférence pour le contrôle de la segmentation 
    return im

def infer_canny(prompt, negative_prompt, image):
    # votre fonction d'inférence pour un contrôle efficace 
    return im

with gr.Blocks(theme='gradio/soft') as demo:
    gr.Markdown("## Stable Diffusion with Different Controls")
    gr.Markdown("In this app, you can find different ControlNets with different filters. ")


    with gr.Tab("ControlNet on Canny Filter "):
        prompt_input_canny = gr.Textbox(label="Prompt")
        negative_prompt_canny = gr.Textbox(label="Negative Prompt")
        canny_input = gr.Image(label="Input Image")
        canny_output = gr.Image(label="Output Image")
        submit_btn = gr.Button(value = "Submit")
        canny_inputs = [prompt_input_canny, negative_prompt_canny, canny_input]
        submit_btn.click(fn=infer_canny, inputs=canny_inputs, outputs=[canny_output])
        
    with gr.Tab("ControlNet with Semantic Segmentation"):
        prompt_input_seg = gr.Textbox(label="Prompt")
        negative_prompt_seg = gr.Textbox(label="Negative Prompt")
        seg_input = gr.Image(label="Image")
        seg_output = gr.Image(label="Output Image")
        submit_btn = gr.Button(value = "Submit")
        seg_inputs = [prompt_input_seg, negative_prompt_seg, seg_input]
        submit_btn.click(fn=infer_segmentation, inputs=seg_inputs, outputs=[seg_output])

demo.launch()

La démo ci-dessus ressemblera à ce qui suit :

Créer notre Space

Une fois notre application écrite, nous pouvons créer un espace Hugging Face pour héberger notre application. Vous pouvez aller sur huggingface.co, cliquer sur votre profil en haut à droite et sélectionner “New Space”.

Nous pouvons nommer notre Space, choisir une licence et sélectionner « Gradio » comme Space SDK.

Après avoir créé le Space, vous pouvez soit utiliser les instructions ci-dessous pour cloner le dépôt localement, ajouter vos fichiers et pousser, soit utiliser l’interface graphique pour créer les fichiers et écrire le code dans le navigateur.

Pour télécharger votre fichier de candidature, cliquez sur « Add File » et faites glisser votre fichier.

Enfin, nous devons créer un fichier appelé requirements.txt et ajouter les conditions requises pour notre projet. Assurez-vous d’installer les versions de jax, diffusers et autres dépendances comme ci-dessous.

-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
jax[cuda11_cudnn805]
jaxlib
git+https://github.com/huggingface/diffusers@main
opencv-python
transformers
flax

Nous vous accorderons une dotation GPU afin que votre application puisse fonctionner sur GPU.

Nous avons un classement hébergé ici et nous distribuerons des prix à partir de ce classement. Pour que votre Space apparaisse sur le leaderboard, éditez simplement README.md de votre Space pour avoir le tag jax-diffusers-event sous les tags comme ci-dessous :

---
title: Canny Coyo1m
emoji: 💜 
...py
tags:
- jax-diffusers-event
---

Prix

Pour ce sprint, nous aurons de nombreux prix. Nous choisirons les dix premiers projets de ce classement, vous devez donc tagger votre Space pour le classement afin que votre soumission soit complète, comme indiqué dans la section ci-dessus. Les projets sont classés en fonction du nombre de j’aimes, nous augmenterons donc partagerons vos Spaces pour en augmenter la visibilité pour que les gens puissent voter en laissant un j’aime sur votre Space. Nous sélectionnerons les dix premiers projets du classement et le jury votera pour déterminer les trois premières places. Ces projets seront mis en valeur par Google et Hugging Face. Les interfaces élaborées ainsi que les projets dont les bases de code et les modèles sont en libre accès augmenteront probablement les chances de gagner des prix.

Les prix sont les suivants et sont remis à chaque membre de l’équipe :

Première place : Un bon d’achat de 150 $ à dépenser sur le Hugging Face Store, un abonnement d’un an à Hugging Face Hub PRO, le livre Natural Language Processing with Transformers.

Deuxième place : Un bon d’achat de 125$ à dépenser sur le Hugging Face Store, un abonnement d’un an à Hugging Face Hub PRO.

Troisième place : Un bon d’achat de 100 $ à dépenser sur le Hugging Face Store, un abonnement d’un an à Hugging Face Hub PRO.

Les dix premiers projets du classement (indépendamment de la décision du jury) gagneront un kit de merch exclusivement conçu pour ce sprint par Hugging Face, ainsi qu’un kit de merch séparé JAX de Google.

Jury

Le jury de ce sprint était composé des personnes suivantes :

1. Robin Rombach, Stability AI
2. Huiwen Chang, Google Research
3. Jun-Yan Zhu, Carnegie Mellon University
4. Merve Noyan, Hugging Face

FAQ

Dans cette section, nous rassemblons les réponses aux questions fréquemment posées sur notre canal discord.

Comment utiliser VSCode avec TPU VM ?

Vous pouvez suivre ce guide général sur la façon d’utiliser VSCode remote pour se connecter à Google Cloud VMs. Une fois que c’est configuré, vous pouvez développer sur la VM TPU en utilisant VSCode.

Pour obtenir votre IP externe, utilisez cette commande :

gcloud compute tpus tpu-vm describe <node_name> --zone=<zone>

Elle devrait être listée sous ‘accessConfig’ -> ‘externalIp’

Comment tester votre code localement ?

Puisque les membres de l’équipe partagent la VM TPU, il peut être pratique d’écrire et de tester votre code localement sur une unité centrale pendant que vos coéquipiers exécutent le processus d’entraînement sur la VM. Pour effectuer des tests locaux, il est important de mettre le drapeau xla_force_host_platform_device_count à 4. Pour en savoir plus, consultez la documentation.

Gagnants du sprint

Les 10 meilleurs projets (basés sur le nombre de likes sur leurs démos) sont disponibles sur ce classement. Nous avons soumis ce classement à notre jury pour qu’il juge les 10 meilleurs projets sur la base de plusieurs facteurs tels que les points de contrôle du modèle, les jeux de données et les bases de code open-source, l’exhaustivité du modèle et des cartes de jeux de données, etc. En conséquence, les trois projets suivants sont sortis vainqueurs :

1. ControlNet pour la décoration intérieure
2. ControlNet pour le réglage de la luminosité
3. Stable Diffusion avec contrôle manuel