Contrôlez les modèles de diffusion en ajoutant des conditions supplémentaires

Auteurs : Sai Bharghav, Prathmesh, Kavya, Alisha, Amritangshu

Les modèles de diffusion texte-image ont fait des progrès impressionnants ces dernières années dans la génération d'images de haute qualité à partir de descriptions textuelles. Ces modèles apprennent à transformer le langage naturel en espace de pixels, ce qui leur permet de générer des images réalistes qui correspondent à la description d'entrée. Cependant, il y a encore place à l'amélioration en termes de contrôle sur les images générées. L'ajout d'un contrôle conditionnel aux modèles de diffusion texte-image peut permettre aux utilisateurs de guider plus précisément la génération d'images, ce qui se traduit par des sorties plus précises et personnalisées. Dans ce contexte, l'utilisation de croquis, de contours et de poses comme moyen de contrôle conditionnel devient de plus en plus populaire. Dans ce blog, nous explorerons le concept de contrôle conditionnel dans les modèles de diffusion texte-image,

Dans ce blog, nous explorerons le concept de combinaison de ControlNet avec une diffusion stable et son potentiel pour améliorer la génération de texte en image.

Qu'est-ce que la diffusion stable ?

La diffusion stable est une variante du modèle de diffusion, un modèle génératif qui apprend à générer des images de haute qualité en modélisant la distribution des sources de bruit intermédiaires. Le modèle de diffusion est un processus séquentiel qui commence par une distribution de bruit fixe et applique une séquence de transformations, chacune modifiant la distribution de bruit pour produire une nouvelle distribution qui correspond mieux à la distribution cible. La distribution finale est utilisée pour générer l'image de sortie. Le modèle de diffusion s'est avéré efficace pour générer des images de haute qualité, mais il peut être difficile de s'entraîner en raison de la nécessité d'estimer des distributions conditionnelles complexes.

La diffusion est le processus qui se déroule à l'intérieur du composant rose "créateur d'informations d'image". Ayant les incorporations de jetons qui représentent le texte d'entrée et un tableau d'informations d'image de départ aléatoire (ceux-ci sont également appelés latents ), le processus produit un tableau d'informations que le décodeur d'image utilise pour peindre l'image finale.

Cette procédure se déroule par étapes. Chaque étape ajoute des données plus pertinentes. Nous pouvons inspecter le réseau de latentes aléatoires et voir comment il se traduit en bruit visuel pour avoir une idée du processus. Dans ce cas, l'inspection visuelle consiste à la faire passer par un décodeur d'image. La diffusion se produit en plusieurs étapes, chaque étape fonctionnant sur un tableau de latentes d'entrée et produisant un autre tableau de latentes qui ressemble plus étroitement au texte d'entrée ainsi qu'à toutes les informations visuelles que le modèle a récupérées à partir de toutes les images sur lesquelles le modèle a été formé. Un ensemble de ces latentes peut être visualisé pour voir quelles informations sont ajoutées à chaque étape.

U sing Diffusion stable sur python à l'aide du pipeline diffuseurs

#!pip install -q diffusers==0.14.0 transformers xformers git+https://github.com/huggingface/accelerate.git
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
import torch

repo_id = "stabilityai/stable-diffusion-2-base"
pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.float16, revision="fp16")

pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

prompt = "Panda in space"
image = pipe(prompt, num_inference_steps=25).images[0]
image.save("astronaut.png")

      
                

                    
                    
                
            

ControlNet est une structure de réseau neuronal qui permet un contrôle explicite sur des modèles de grande diffusion pré-entraînés, ce qui permet de prendre en charge des conditions d'entrée supplémentaires. Le ControlNet est conçu pour apprendre des conditions spécifiques à une tâche de bout en bout, ce qui rend le processus d'apprentissage robuste même lorsque l'ensemble de données d'apprentissage est petit, avec moins de 50 000 échantillons. Cette capacité fait de ControlNet un outil puissant pour la génération de texte en image, où les variables de contrôle peuvent être utilisées pour spécifier les attributs souhaités des images générées, telles que la forme, la couleur ou la texture.

Comment ça fonctionne

ControlNet fonctionne en copiant les poids des blocs de réseau neuronal en deux copies : une copie « verrouillée » et une copie « entraînable ». La copie "verrouillée" préserve le modèle d'origine, tandis que la copie "entraînable" est utilisée pour apprendre les conditions spécifiées.

Pendant la formation, ControlNet utilise une technique appelée « convolution zéro », qui est une convolution 1x1 avec à la fois un poids et un biais initialisés comme des zéros. Avant l'entraînement, toutes les convolutions nulles produisent des zéros et ControlNet ne cause aucune distorsion au modèle d'origine. La copie "entraînable" apprend les conditions spécifiées, tandis que la copie "verrouillée" reste inchangée.

Cette approche a de nombreux avantages. Premièrement, il permet de s'entraîner sur de petits ensembles de données de paires d'images sans compromettre les modèles de diffusion prêts à la production. Deuxièmement, aucune couche n'est formée à partir de zéro, ce qui signifie que le modèle d'origine est sûr. Troisièmement, cette approche permet une formation sur des appareils à petite échelle ou même personnels, ce qui la rend pratique pour un large éventail d'applications.

Utilisation du pipeline StableDiffusion ControlNet

nous utilisons un pipeline pour la génération de texte en image à l'aide de Stable Diffusion avec les conseils de ControlNet

Modèles pré-formés disponibles :

• Contrôle à l'aide de la détection des bords : lllyasviel/sd-controlnet-canny : formé avec la détection des bords astucieux, utilisé pour contrôler les bords de l'image.
• Contrôle à l'aide de la détection de pose : lllyasviel/sd-controlnet_openpose : formé avec l'image d'os OpenPose, utilisé pour contrôler les poses humaines.
• Contrôle à l'aide de la détection de gribouillis : lllyasviel/sd-controlnet_scribble : formé avec des gribouillis humains, utilisé pour contrôler les contours de l'image.
• lllyasviel/sd-controlnet-depth : formé avec l'estimation de profondeur Midas, utilisé pour contrôler la profondeur de l'image.
• lllyasviel/sd-controlnet-hed : formé avec la détection de bord HED (bord doux), utilisé pour contrôler les bords doux de l'image.
• lllyasviel/sd-controlnet-mlsd : formé avec la détection de ligne M-LSD, utilisé pour contrôler les lignes droites dans une image.
• lllyasviel/sd-controlnet-normal : formé avec la carte normale, utilisé pour contrôler les normales d'image.
• lllyasviel/sd-controlnet_seg : formé avec la segmentation sémantique, utilisé pour contrôler la segmentation sémantique des images.

Tout d'abord, nous devons configurer notre environnement et importer une diffusion stable à partir du visage étreignant, ici nous utilisons stable-diffusion-v1–5

Le code installe certains packages à l'aide de pip, puis importe les modules et les classes de ces packages.

La première ligne installe les packages suivants :

• diffusers==0.14.0: Un package pour implémenter des modèles de diffusion dans PyTorch.
• transformers: un package pour les tâches de traitement du langage naturel (TAL) telles que la classification de texte et la génération de langage.
• xformers: Un package pour les modèles de transformateurs, qui sont un type d'architecture de réseau neuronal utilisé dans les tâches NLP.
• git+https://github.com/huggingface/accelerate.git: Un package pour optimiser la formation des modèles d'apprentissage en profondeur.

!pip install -q diffusers==0.14.0 transformers xformers git+https://github.com/huggingface/accelerate.git
!pip install -q opencv-contrib-python
!pip install -q controlnet_aux

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
import torch

controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
)

from diffusers import StableDiffusionControlNetPipeline
from diffusers.utils import load_image

image = load_image(
    "/content/daniel-craig-007.jpg-303a730.png"
)
image

import cv2
from PIL import Image
import numpy as np

image = np.array(image)

low_threshold = 100
high_threshold = 200

image = cv2.Canny(image, low_threshold, high_threshold)
image = image[:, :, None]
image = np.concatenate([image, image, image], axis=2)
canny_image = Image.fromarray(image)
canny_image

      
                

                    
                    
                
            

      
                

                    
                    
                
            

La sortie de l' pipeobjet est un ensemble de quatre images qui sont générées en fonction des invites et des paramètres de contrôle fournis. Les invites sont fournies sous la forme d'une liste de quatre chaînes, chaque chaîne étant un nom de célébrité suivi d'un message indiquant les qualités d'image souhaitées. Les messages d'invite indiquent un désir d'images de "meilleure qualité" et "extrêmement détaillées". Le negative_promptparamètre est défini pour indiquer les qualités d'image souhaitées de manière négative en demandant des images "monochrome", "lowres", "mauvaise anatomie", "pire qualité" et "basse qualité". Le generatorparamètre est utilisé pour contrôler le caractère aléatoire des images générées, et le num_inference_stepsparamètre définit le nombre d'itérations utilisées pour générer chaque image.

def image_grid(imgs, rows, cols):
    assert len(imgs) == rows * cols

    w, h = imgs[0].size
    grid = Image.new("RGB", size=(cols * w, rows * h))
    grid_w, grid_h = grid.size

    for i, img in enumerate(imgs):
        grid.paste(img, box=(i % cols * w, i // cols * h))
    return grid

prompt = ", best quality, extremely detailed"
prompt = [t + prompt for t in ["Tom Cruise", "Donald Trump", "rihanna", "taylor swift"]]
generator = [torch.Generator(device="cpu").manual_seed(2) for i in range(len(prompt))]

output = pipe(
    prompt,
    canny_image,
    negative_prompt=["monochrome, lowres, bad anatomy, worst quality, low quality"] * len(prompt),
    generator=generator,
    num_inference_steps=20,
)

image_grid(output.images, 2, 2)

      
                

                    
                    
                
            

Ce code charge un OpenposeDetectormodèle pré-formé à partir du hub de modèles Hugging Face à l'aide de la from_pretrained()méthode et l'applique à une image pour détecter les poses humaines. Il charge également un pré-formé ControlNetModelet un StableDiffusionControlNetPipelineà partir du diffuserspackage, définit certains paramètres de contrôle pour le pipeline et permet certaines optimisations de la mémoire et des performances.

from controlnet_aux import OpenposeDetector
model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
poses = model(image)
#image_grid(poses, 2, 2)
poses
controlnet = ControlNetModel.from_pretrained(
    "fusing/stable-diffusion-v1-5-controlnet-openpose", torch_dtype=torch.float16
)
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    model_id,
    controlnet=controlnet,
    torch_dtype=torch.float16,
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()
pipe.enable_xformers_memory_efficient_attention()

      
                

                    
                    
                
            

      
                

                    
                    
                
            

      
                

                    
                    
                
            

Nous voulions voir comment former un modèle ControlNet sur un jeu de données. Donc, pour atteindre notre objectif de contrôler SD, nous avons pris la tâche de remplir les cercles de couleurs. Notre approche consistera à utiliser une tâche simple pour démontrer comment ControlNet peut apprendre des conditions spécifiques à une tâche de bout en bout. Nous utiliserons une invite décrivant notre cible et incluant une « Image de contrôle (image source) ».

La diffusion stable a déjà été formée sur des milliards d'images, et elle sait déjà ce que sont les couleurs et les formes. Cependant, il ne comprend pas la signification de « Image de contrôle (image source) » dans notre invite. En ajoutant ControlNet, nous pouvons apprendre à SD à comprendre cette condition d'entrée et à remplir le cercle avec la couleur spécifiée.

L'ensemble de données

L'ensemble de données "fill50k" contient une collection de 50 000 paires d'images de cercles, qui sont stockées dans deux dossiers : "source" et "cible". Le dossier "source" contient des images de cercles avec des lignes uniquement, tandis que le dossier "cible" contient des images de cercles remplis d'une couleur spécifique.

En plus des images, l'ensemble de données comprend un fichier "prompt.json", qui fournit des informations sur chaque image. Chaque invite se présente sous la forme d'une description du cercle et de la couleur de l'arrière-plan. Par exemple, une invite peut lire "un cercle bleu sur fond jaune".

Pour utiliser cet ensemble de données avec PyTorch, vous devrez écrire un script qui lit les images et leurs invites correspondantes à partir du dossier "fill50k" et les prépare pour la formation. Cela impliquera de charger les images dans des tenseurs PyTorch, de créer des lots de données et d'effectuer tout prétraitement nécessaire, tel que la normalisation ou l'augmentation des données.

Dans l'ensemble, l'ensemble de données "fill50k" est conçu pour être utilisé pour former des modèles d'apprentissage automatique afin de remplir des cercles avec des couleurs spécifiques en fonction des invites fournies dans l'ensemble de données. Il s'agit d'une ressource utile pour les chercheurs et les développeurs travaillant sur la génération d'images et les tâches connexes à l'aide de techniques d'apprentissage en profondeur.

Chargement du jeu de données

import json
import cv2
import numpy as np

from torch.utils.data import Dataset


class MyDataset(Dataset):
    def __init__(self):
        self.data = []
        with open('./training/fill50k/prompt.json', 'rt') as f:
            for line in f:
                self.data.append(json.loads(line))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]

        source_filename = item['source']
        target_filename = item['target']
        prompt = item['prompt']

        source = cv2.imread('./training/fill50k/' + source_filename)
        target = cv2.imread('./training/fill50k/' + target_filename)

        # Do not forget that OpenCV read images in BGR order.
        source = cv2.cvtColor(source, cv2.COLOR_BGR2RGB)
        target = cv2.cvtColor(target, cv2.COLOR_BGR2RGB)

        # Normalize source images to [0, 1].
        source = source.astype(np.float32) / 255.0

        # Normalize target images to [-1, 1].
        target = (target.astype(np.float32) / 127.5) - 1.0

        return dict(jpg=target, txt=prompt, hint=source)

Notez que tous les poids à l'intérieur du ControlNet sont également copiés à partir de SD afin qu'aucune couche ne soit formée à partir de zéro, et vous continuez à affiner l'ensemble du modèle.

Former le modèle

Ce script Python utilise PyTorch Lightning pour former un modèle d'apprentissage automatique personnalisé sur un ensemble de données personnalisé. Il charge le modèle pré-formé, définit les configurations telles que la taille du lot, la fréquence de l'enregistreur, le taux d'apprentissage, etc. Il crée ensuite un DataLoader pour charger les données de l'ensemble de données personnalisé, crée un enregistreur d'images pour enregistrer la progression de la formation et configure le PyTorch Lightning Trainer pour gérer le processus de formation. Enfin, il entraîne le modèle à l'aide de Trainer et de DataLoader, et enregistre la progression à l'aide de ImageLogger.

import pytorch_lightning as pl
from torch.utils.data import DataLoader
from tutorial_dataset import MyDataset
from cldm.logger import ImageLogger
from cldm.model import create_model, load_state_dict


# Configs
resume_path = './models/control_sd15_ini.ckpt'
batch_size = 4
logger_freq = 300
learning_rate = 1e-5
sd_locked = True
only_mid_control = False


# First use cpu to load models. Pytorch Lightning will automatically move it to GPUs.
model = create_model('./models/cldm_v15.yaml').cpu()
model.load_state_dict(load_state_dict(resume_path, location='cpu'))
model.learning_rate = learning_rate
model.sd_locked = sd_locked
model.only_mid_control = only_mid_control


# Misc
dataset = MyDataset()
dataloader = DataLoader(dataset, num_workers=0, batch_size=batch_size, shuffle=True)
logger = ImageLogger(batch_frequency=logger_freq)
trainer = pl.Trainer(gpus=1, precision=32, callbacks=[logger])


# Train!
trainer.fit(model, dataloader)

      
                

                    
                    
                
            

      
                

                    
                    
                
            

      
                

                    
                    
                
            

Dans l'ensemble, le processus de génération d'une image à partir d'un signal d'invite et de contrôle implique la formation d'un modèle d'apprentissage en profondeur sur un ensemble de données de paires entrée-sortie et l'ajustement des paramètres du modèle afin qu'il apprenne à combiner les entrées pour générer les sorties.

Apprentissages

Préparation des données et du modèle :

• Les modèles pré-formés peuvent être utiles pour une variété de tâches et peuvent être facilement intégrés dans des projets.
• Différents modèles pré-formés sont disponibles pour contrôler différents aspects d'une image, tels que les bords, la profondeur et les poses humaines.
• Les bibliothèques open source telles que PyTorch peuvent être utilisées pour la formation et le déploiement de modèles d'apprentissage automatique.

• La diffusion stable est une technique puissante pour générer des images de haute qualité basées sur des entrées de contrôle.
• L'utilisation d'un modèle de diffusion avec un réseau de contrôle peut permettre un contrôle plus précis de l'image générée.
• Plusieurs invites peuvent être utilisées pour générer différentes images basées sur la même entrée de contrôle.

• Des techniques de traitement d'image telles que la détection des contours Canny et OpenPose peuvent être utilisées pour extraire des informations spécifiques des images, qui peuvent ensuite être utilisées comme entrées de contrôle pour la génération d'images.
• Différents types de techniques de traitement d'image conviennent à différents types d'entrées de contrôle, telles que les bords, les poses et les contours.

Sources

Nous avons utilisé les ressources suivantes dans notre analyse et notre base de code à titre de référence.