Steuern Sie Diffusionsmodelle, indem Sie zusätzliche Bedingungen hinzufügen

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

Text-zu-Bild-Diffusionsmodelle haben in den letzten Jahren beeindruckende Fortschritte bei der Generierung hochwertiger Bilder aus Textbeschreibungen gemacht. Diese Modelle lernen, natürliche Sprache in Pixelraum umzuwandeln, sodass sie realistische Bilder erzeugen können, die der Eingabebeschreibung entsprechen. Hinsichtlich der Kontrolle über die erzeugten Bilder besteht jedoch noch Verbesserungspotenzial. Das Hinzufügen einer bedingten Steuerung zu Text-zu-Bild-Diffusionsmodellen kann es Benutzern ermöglichen, die Bilderzeugung präziser zu steuern, was zu genaueren und maßgeschneiderten Ergebnissen führt. In diesem Zusammenhang erfreut sich der Einsatz von Skizzen, Umrissen und Posen als Mittel der bedingten Kontrolle immer größerer Beliebtheit. In diesem Blog werden wir das Konzept der bedingten Kontrolle in Text-zu-Bild-Diffusionsmodellen untersuchen.

In diesem Blog werden wir das Konzept der Kombination von ControlNet mit stabiler Diffusion und sein Potenzial zur Verbesserung der Text-zu-Bild-Generierung untersuchen.

Was ist stabile Diffusion?

Stabile Diffusion ist eine Variante des Diffusionsmodells, eines generativen Modells, das lernt, qualitativ hochwertige Bilder zu erzeugen, indem es die Verteilung zwischenliegender Rauschquellen modelliert. Das Diffusionsmodell ist ein sequenzieller Prozess, der mit einer festen Rauschverteilung beginnt und eine Folge von Transformationen anwendet, von denen jede die Rauschverteilung modifiziert, um eine neue Verteilung zu erzeugen, die besser zur Zielverteilung passt. Die endgültige Verteilung wird verwendet, um das Ausgabebild zu generieren. Das Diffusionsmodell hat sich bei der Generierung hochwertiger Bilder als wirksam erwiesen, das Training kann jedoch schwierig sein, da komplexe bedingte Verteilungen geschätzt werden müssen.

Diffusion ist der Prozess, der innerhalb der rosa Komponente „Bildinformationsersteller“ stattfindet. Mit den Token-Einbettungen, die den Eingabetext darstellen, und einem zufälligen Startbild- Informationsarray (diese werden auch Latents genannt ) erzeugt der Prozess ein Informationsarray, das der Bilddecoder zum Zeichnen des endgültigen Bildes verwendet.

Dieser Vorgang wird stufenweise durchgeführt. Mit jedem Schritt werden weitere relevante Daten hinzugefügt. Wir können das zufällige Latent-Array untersuchen und sehen, wie es sich in visuelles Rauschen umsetzt, um ein Gefühl für den Prozess zu bekommen. In diesem Fall erfordert die visuelle Prüfung, dass das Bild durch einen Bilddecoder geleitet wird. Die Diffusion erfolgt in mehreren Schritten, wobei jeder Schritt ein Eingabe-Latent-Array bearbeitet und ein weiteres Latent-Array erzeugt, das dem Eingabetext sowie allen visuellen Informationen, die das Modell aus allen Bildern, auf denen das Modell trainiert wurde, besser ähnelt, ähnlicher ist. Eine Reihe dieser latenten Daten kann visualisiert werden, um zu sehen, welche Informationen bei jedem Schritt hinzugefügt werden.

Verwendung einer stabilen Diffusion auf Python mithilfe der Diffusor-Pipeline

#!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 ist eine neuronale Netzwerkstruktur, die eine explizite Steuerung vorab trainierter Modelle mit großer Diffusion ermöglicht und so die Unterstützung zusätzlicher Eingabebedingungen ermöglicht. Das ControlNet ist darauf ausgelegt, aufgabenspezifische Bedingungen durchgängig zu lernen, wodurch der Lernprozess robust ist, selbst wenn der Trainingsdatensatz klein ist und weniger als 50.000 Stichproben umfasst. Diese Fähigkeit macht ControlNet zu einem leistungsstarken Tool für die Text-zu-Bild-Generierung, bei dem die Steuervariablen verwendet werden können, um gewünschte Attribute der generierten Bilder wie Form, Farbe oder Textur festzulegen.

Wie es funktioniert

ControlNet kopiert die Gewichte neuronaler Netzwerkblöcke in zwei Kopien: eine „gesperrte“ Kopie und eine „trainierbare“ Kopie. Die „gesperrte“ Kopie behält das Originalmodell bei, während die „trainierbare“ Kopie zum Erlernen der angegebenen Bedingungen verwendet wird.

Während des Trainings verwendet ControlNet eine Technik namens „Null-Faltung“, bei der es sich um eine 1x1-Faltung handelt, bei der sowohl Gewichtung als auch Bias als Nullen initialisiert werden. Vor dem Training geben alle Nullfaltungen Nullen aus und ControlNet verursacht keine Verzerrung des Originalmodells. Die „trainierbare“ Kopie lernt die festgelegten Bedingungen, während die „gesperrte“ Kopie unverändert bleibt.

Dieser Ansatz hat mehrere Vorteile. Erstens ermöglicht es das Training an kleinen Datensätzen von Bildpaaren, ohne die produktionsbereiten Diffusionsmodelle zu beeinträchtigen. Zweitens wird keine Ebene von Grund auf trainiert, was bedeutet, dass das Originalmodell sicher ist. Drittens ermöglicht dieser Ansatz das Training auf kleinen oder sogar persönlichen Geräten, was ihn für eine Vielzahl von Anwendungen praktisch macht.

Verwendung der StableDiffusion ControlNet-Pipeline

Wir verwenden eine Pipeline für die Text-zu-Bild-Generierung unter Verwendung von Stable Diffusion mit ControlNet-Anleitung

Vortrainierte Modelle verfügbar:

• Steuerung mithilfe der Kantenerkennung : lllyasviel/sd-controlnet-canny: Trainiert mit der Kantenerkennung, die zur Steuerung von Bildkanten verwendet wird.
• Steuerung mithilfe der Posenerkennung: lllyasviel/sd-controlnet_openpose: Trainiert mit dem OpenPose-Knochenbild, das zur Steuerung menschlicher Posen verwendet wird.
• Steuerung mithilfe der Scribble-Erkennung: lllyasviel/sd-controlnet_scribble: Mit menschlichen Kritzeleien trainiert, wird zur Steuerung von Bildkonturen verwendet.
• lllyasviel/sd-controlnet- Depth: Trainiert mit Midas-Tiefenschätzung, wird zur Steuerung der Bildtiefe verwendet.
• lllyasviel/sd-controlnet-hed: Mit HED-Kantenerkennung (weiche Kante) trainiert, die zur Steuerung weicher Bildkanten verwendet wird.
• lllyasviel/sd-controlnet-mlsd: Mit M-LSD-Linienerkennung trainiert, wird zur Steuerung gerader Linien in einem Bild verwendet.
• lllyasviel/sd-controlnet-normal: Mit Normalenkarte trainiert, wird zur Steuerung von Bildnormalen verwendet.
• lllyasviel/sd-controlnet_seg: Mit semantischer Segmentierung trainiert, wird zur Steuerung der semantischen Bildsegmentierung verwendet.

Zuerst müssen wir unsere Umgebung einrichten und eine stabile Diffusion aus dem Hugging Face importieren. Hier verwenden wir Stable-Diffusion-v1–5

Der Code installiert einige Pakete mit pip und importiert dann Module und Klassen aus diesen Paketen.

Die erste Zeile installiert die folgenden Pakete:

• diffusers==0.14.0: Ein Paket zur Implementierung von Diffusionsmodellen in PyTorch.
• transformers: Ein Paket für Aufgaben der Verarbeitung natürlicher Sprache (NLP), wie z. B. Textklassifizierung und Sprachgenerierung.
• xformers: Ein Paket für Transformatormodelle, eine Art neuronale Netzwerkarchitektur, die in NLP-Aufgaben verwendet wird.
• git+https://github.com/huggingface/accelerate.git: Ein Paket zur Optimierung des Trainings von Deep-Learning-Modellen.

!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

      
                

                    
                    
                
            

      
                

                    
                    
                
            

Die Ausgabe des pipeObjekts ist ein Satz von vier Bildern, die basierend auf den Eingabeaufforderungen und den bereitgestellten Steuerparametern generiert werden. Die Eingabeaufforderungen werden als Liste mit vier Zeichenfolgen bereitgestellt, wobei jede Zeichenfolge aus einem Prominentennamen gefolgt von einer Nachricht besteht, die die gewünschten Bildqualitäten angibt. Die Aufforderungsmeldungen deuten auf den Wunsch nach „bester Qualität“ und „extrem detaillierten“ Bildern hin. Der negative_promptParameter ist so eingestellt, dass die gewünschten Bildqualitäten negativ angezeigt werden, indem Bilder mit der Aufforderung „Monochrom“, „Niedrige Auflösung“, „schlechte Anatomie“, „schlechteste Qualität“ und „Geringe Qualität“ angefordert werden. Der generatorParameter wird verwendet, um die Zufälligkeit der generierten Bilder zu steuern, und der num_inference_stepsParameter legt die Anzahl der Iterationen fest, die zum Generieren jedes Bildes verwendet werden.

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)

      
                

                    
                    
                
            

Dieser Code lädt mithilfe der Methode ein vorab trainiertes OpenposeDetectorModell aus dem Hugging Face-Modell-Hub und wendet es auf ein Bild an, um menschliche Posen zu erkennen. from_pretrained()Außerdem lädt es ein vorab trainiertes ControlNetModelund ein StableDiffusionControlNetPipelineaus dem diffusersPaket, legt einige Steuerparameter für die Pipeline fest und ermöglicht einige Speicher- und Leistungsoptimierungen.

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()

      
                

                    
                    
                
            

      
                

                    
                    
                
            

      
                

                    
                    
                
            

Wir wollten sehen, wie wir ein ControlNet-Modell anhand eines Datensatzes trainieren können. Um unser Ziel, SD zu kontrollieren, zu erreichen, haben wir es uns zur Aufgabe gemacht, Kreise mit Farben zu füllen. Unser Ansatz besteht darin, anhand einer einfachen Aufgabe zu demonstrieren, wie ControlNet aufgabenspezifische Bedingungen durchgängig lernen kann. Wir verwenden eine Eingabeaufforderung, die unser Ziel beschreibt und ein „Kontrollbild (Quellbild)“ enthält.

Die stabile Diffusion wurde bereits auf Milliarden von Bildern trainiert und weiß bereits, welche Farben und Formen es gibt. Es versteht jedoch nicht die Bedeutung des „Kontrollbilds (Quellbild)“ in unserer Eingabeaufforderung. Durch das Hinzufügen von ControlNet können wir SD beibringen, diese Eingabebedingung zu verstehen und den Kreis mit der angegebenen Farbe zu füllen.

Der Datensatz

Der Datensatz „fill50k“ enthält eine Sammlung von 50.000 Bildpaaren von Kreisen, die in zwei Ordnern gespeichert sind: „Quelle“ und „Ziel“. Der Ordner „Quelle“ enthält nur Bilder von Kreisen mit Linien, während der Ordner „Ziel“ Bilder von Kreisen enthält, die mit einer bestimmten Farbe gefüllt sind.

Zusätzlich zu den Bildern enthält der Datensatz eine „prompt.json“-Datei, die Informationen zu jedem Bild bereitstellt. Jede Eingabeaufforderung besteht aus einer Beschreibung des Kreises und der Farbe des Hintergrunds. Eine Aufforderung könnte beispielsweise „ein blauer Kreis auf gelbem Hintergrund“ lauten.

Um diesen Datensatz mit PyTorch zu verwenden, müssen Sie ein Skript schreiben, das die Bilder und die entsprechenden Eingabeaufforderungen aus dem Ordner „fill50k“ liest und sie für das Training vorbereitet. Dazu gehört das Laden der Bilder in PyTorch-Tensoren, das Erstellen von Datenstapeln und die Durchführung aller notwendigen Vorverarbeitungen wie Normalisierung oder Datenerweiterung.

Insgesamt ist der „fill50k“-Datensatz darauf ausgelegt, zum Trainieren von Modellen für maschinelles Lernen verwendet zu werden, um Kreise basierend auf den im Datensatz bereitgestellten Eingabeaufforderungen mit bestimmten Farben zu füllen. Es ist eine nützliche Ressource für Forscher und Entwickler, die mithilfe von Deep-Learning-Techniken an der Bildgenerierung und damit verbundenen Aufgaben arbeiten.

Laden des Datensatzes

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)

Beachten Sie, dass alle Gewichte im ControlNet auch von SD kopiert werden, sodass keine Ebene von Grund auf trainiert wird und Sie immer noch das gesamte Modell optimieren.

Trainieren des Modells

Dieses Python-Skript verwendet PyTorch Lightning, um ein benutzerdefiniertes Modell für maschinelles Lernen auf einem benutzerdefinierten Datensatz zu trainieren. Es lädt das vorab trainierte Modell und legt die Konfigurationen wie Batch-Größe, Logger-Frequenz, Lernrate und mehr fest. Anschließend erstellt es einen DataLoader zum Laden von Daten aus dem benutzerdefinierten Datensatz, erstellt einen Bildlogger zum Aufzeichnen des Trainingsfortschritts und richtet den PyTorch Lightning Trainer für die Abwicklung des Trainingsprozesses ein. Schließlich trainiert es das Modell mit dem Trainer und DataLoader und protokolliert den Fortschritt mit dem 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)

      
                

                    
                    
                
            

      
                

                    
                    
                
            

      
                

                    
                    
                
            

Insgesamt umfasst der Prozess der Generierung eines Bildes aus einem Eingabeaufforderungs- und Steuersignal das Trainieren eines Deep-Learning-Modells anhand eines Datensatzes von Eingabe-Ausgabe-Paaren und das Anpassen der Parameter des Modells, damit es lernt, die Eingaben zu kombinieren, um die Ausgaben zu generieren.

Erkenntnisse

Daten- und Modellvorbereitung:

• Vorab trainierte Modelle können für eine Vielzahl von Aufgaben nützlich sein und leicht in Projekte integriert werden.
• Für die Steuerung verschiedener Aspekte eines Bildes, wie Kanten, Tiefe und menschliche Posen, stehen verschiedene vorab trainierte Modelle zur Verfügung.
• Open-Source-Bibliotheken wie PyTorch können zum Trainieren und Bereitstellen von Modellen für maschinelles Lernen verwendet werden.

• Stabile Diffusion ist eine leistungsstarke Technik zur Erzeugung hochwertiger Bilder basierend auf Steuereingaben.
• Die Verwendung eines Diffusionsmodells mit einem Kontrollnetzwerk kann eine genauere Kontrolle über das erzeugte Bild ermöglichen.
• Mehrere Eingabeaufforderungen können verwendet werden, um unterschiedliche Bilder basierend auf derselben Steuereingabe zu generieren.

• Bildverarbeitungstechniken wie Canny Edge Detection und OpenPose können verwendet werden, um spezifische Informationen aus Bildern zu extrahieren, die dann als Steuereingaben für die Bilderzeugung verwendet werden können.
• Verschiedene Arten von Bildverarbeitungstechniken eignen sich für unterschiedliche Arten von Steuereingaben, z. B. Kanten, Posen und Umrisse.

Quellen

Wir haben die folgenden Ressourcen in unserer Analyse und Codebasis als Referenz verwendet.