Modelos de difusión de control agregando condiciones adicionales

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

Los modelos de difusión de texto a imagen han logrado avances impresionantes en los últimos años en la generación de imágenes de alta calidad a partir de descripciones textuales. Estos modelos aprenden a transformar el lenguaje natural en espacio de píxeles, lo que les permite generar imágenes realistas que coinciden con la descripción de entrada. Sin embargo, todavía hay margen de mejora en términos de control sobre las imágenes generadas. Agregar control condicional a los modelos de difusión de texto a imagen puede permitir a los usuarios guiar la generación de imágenes con mayor precisión, lo que da como resultado resultados más precisos y personalizados. En este contexto, el uso de bocetos, esquemas y poses como medio de control condicional se está volviendo cada vez más popular. En este blog, exploraremos el concepto de control condicional en los modelos de difusión de texto a imagen,

En este blog, exploraremos el concepto de combinar ControlNet con difusión estable y su potencial para mejorar la generación de texto a imagen.

¿Qué es la difusión estable?

La difusión estable es una variante del modelo de difusión, un modelo generativo que aprende a generar imágenes de alta calidad modelando la distribución de fuentes de ruido intermedias. El modelo de difusión es un proceso secuencial que comienza con una distribución de ruido fija y aplica una secuencia de transformaciones, cada una de las cuales modifica la distribución de ruido para producir una nueva distribución que coincida mejor con la distribución objetivo. La distribución final se utiliza para generar la imagen de salida. Se ha demostrado que el modelo de difusión es efectivo para generar imágenes de alta calidad, pero puede ser un desafío entrenarlo debido a la necesidad de estimar distribuciones condicionales complejas.

La difusión es el proceso que tiene lugar dentro del componente rosa "creador de información de imagen". Con las incrustaciones de tokens que representan el texto de entrada y una matriz de información de imagen inicial aleatoria (también denominadas latentes ), el proceso produce una matriz de información que el decodificador de imágenes utiliza para pintar la imagen final.

Este procedimiento se lleva a cabo por etapas. Cada paso agrega más datos relevantes. Podemos inspeccionar la matriz de latentes aleatorias y ver cómo se traduce en ruido visual para tener una idea del proceso. En este caso, la inspección visual implica pasarlo por un decodificador de imágenes. La difusión se produce en varios pasos, con cada paso operando en una matriz de latentes de entrada y produciendo otra matriz de latentes que se parece más al texto de entrada, así como a toda la información visual que el modelo recogió de todas las imágenes en las que se entrenó el modelo. Se puede visualizar un conjunto de estas latentes para ver qué información se agrega en cada paso.

U sando difusión estable en python usando la canalización de difusores

#!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 es una estructura de red neuronal que permite un control explícito sobre modelos de gran difusión preentrenados, lo que permite admitir condiciones de entrada adicionales. ControlNet está diseñado para aprender condiciones específicas de tareas de forma integral, lo que hace que el proceso de aprendizaje sea sólido incluso cuando el conjunto de datos de entrenamiento es pequeño, con menos de 50 000 muestras. Esta capacidad convierte a ControlNet en una poderosa herramienta para la generación de texto a imagen, donde las variables de control se pueden usar para especificar los atributos deseados de las imágenes generadas, como la forma, el color o la textura.

Cómo funciona

ControlNet funciona copiando los pesos de los bloques de redes neuronales en dos copias: una copia "bloqueada" y una copia "entrenable". La copia "bloqueada" conserva el modelo original, mientras que la copia "entrenable" se usa para aprender las condiciones especificadas.

Durante el entrenamiento, ControlNet usa una técnica llamada "convolución cero", que es una convolución 1x1 con el peso y el sesgo inicializados como ceros. Antes del entrenamiento, todas las convoluciones cero generan ceros y ControlNet no causa ninguna distorsión en el modelo original. La copia "entrenable" aprende las condiciones especificadas, mientras que la copia "bloqueada" permanece sin cambios.

Este enfoque tiene varias ventajas. Primero, permite el entrenamiento en pequeños conjuntos de datos de pares de imágenes sin comprometer los modelos de difusión listos para producción. En segundo lugar, ninguna capa se entrena desde cero, lo que significa que el modelo original es seguro. En tercer lugar, este enfoque permite la capacitación en dispositivos a pequeña escala o incluso personales, lo que lo hace práctico para una amplia gama de aplicaciones.

Uso de la canalización StableDiffusion ControlNet

estamos usando un Pipeline para la generación de texto a imagen usando Stable Diffusion con guía ControlNet

Modelos preentrenados disponibles:

• Control mediante detección de bordes : lllyasviel/sd-controlnet-canny: Entrenado con detección de bordes canny, utilizado para controlar los bordes de la imagen.
• Control usando Detección de pose: lllyasviel/sd-controlnet_openpose: Entrenado con imagen ósea OpenPose, utilizado para controlar poses humanas.
• Control mediante detección de garabatos: lllyasviel/sd-controlnet_scribble: entrenado con garabatos humanos, utilizado para controlar los contornos de las imágenes.
• lllyasviel/sd-controlnet- depth: Entrenado con estimación de profundidad Midas, utilizado para controlar la profundidad de la imagen.
• lllyasviel/sd-controlnet-hed: capacitado con detección de bordes HED (borde suave), utilizado para controlar los bordes suaves de la imagen.
• lllyasviel/sd-controlnet-mlsd: Entrenado con detección de línea M-LSD, utilizado para controlar líneas rectas en una imagen.
• lllyasviel/sd-controlnet-normal: Entrenado con mapa normal, usado para controlar normales de imagen.
• lllyasviel/sd-controlnet_seg: Entrenado con segmentación semántica, utilizado para controlar la segmentación semántica de imágenes.

Primero, tenemos que configurar nuestro entorno e importar la difusión estable de la cara abrazada, aquí estamos usando la difusión estable-v1–5

El código instala algunos paquetes usando pip y luego importa módulos y clases de estos paquetes.

La primera línea instala los siguientes paquetes:

• diffusers==0.14.0: Un paquete para implementar modelos de difusión en PyTorch.
• transformers: un paquete para tareas de procesamiento de lenguaje natural (NLP), como clasificación de texto y generación de lenguaje.
• xformers: un paquete para modelos de transformadores, que son un tipo de arquitectura de red neuronal utilizada en tareas de NLP.
• git+https://github.com/huggingface/accelerate.git: un paquete para optimizar el entrenamiento de modelos de aprendizaje profundo.

!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 salida del pipeobjeto es un conjunto de cuatro imágenes que se generan en función de las indicaciones y los parámetros de control proporcionados. Las indicaciones se proporcionan como una lista de cuatro cadenas, cada una de las cuales es el nombre de una celebridad seguida de un mensaje que indica las cualidades de imagen deseadas. Los mensajes de aviso indican un deseo de imágenes de "mejor calidad" y "extremadamente detalladas". El negative_promptparámetro se configura para indicar las calidades de imagen deseadas de forma negativa solicitando imágenes "monocromáticas", "de baja resolución", "mala anatomía", "peor calidad" y "baja calidad". El generatorparámetro se utiliza para controlar la aleatoriedad de las imágenes generadas y el num_inference_stepsparámetro establece el número de iteraciones utilizadas para generar cada imagen.

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)

      
                

                    
                    
                
            

Este código carga un modelo previamente entrenado OpenposeDetectordesde el concentrador de modelos Hugging Face utilizando el from_pretrained()método y lo aplica a una imagen para detectar poses humanas. También carga un preentrenado ControlNetModely un StableDiffusionControlNetPipelinedel diffuserspaquete, establece algunos parámetros de control para la canalización y habilita algunas optimizaciones de memoria y rendimiento.

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

      
                

                    
                    
                
            

      
                

                    
                    
                
            

      
                

                    
                    
                
            

Queríamos ver cómo podemos entrenar un modelo de ControlNet en un conjunto de datos. Entonces, para lograr nuestro objetivo de controlar SD, asumimos la tarea de llenar círculos con colores. Nuestro enfoque será utilizar una tarea simple para demostrar cómo ControlNet puede aprender las condiciones específicas de la tarea de forma integral. Usaremos un aviso que describa nuestro objetivo e incluya una "Imagen de control (Imagen de origen)".

La difusión estable ya se ha entrenado en miles de millones de imágenes y ya sabe qué colores y formas son. Sin embargo, no comprende el significado de la "Imagen de control (Imagen de origen)" en nuestro aviso. Al agregar ControlNet, podemos enseñarle a SD a comprender esta condición de entrada y llenar el círculo con el color especificado.

El conjunto de datos

El conjunto de datos "fill50k" contiene una colección de 50 000 pares de imágenes de círculos, que se almacenan en dos carpetas: "origen" y "objetivo". La carpeta "origen" contiene imágenes de círculos con líneas únicamente, mientras que la carpeta "destino" contiene imágenes de círculos rellenos con un color específico.

Además de las imágenes, el conjunto de datos incluye un archivo "prompt.json", que proporciona información sobre cada imagen. Cada mensaje tiene la forma de una descripción del círculo y el color del fondo. Por ejemplo, un aviso puede decir "un círculo azul sobre un fondo amarillo".

Para usar este conjunto de datos con PyTorch, deberá escribir un script que lea las imágenes y sus indicaciones correspondientes de la carpeta "fill50k" y las prepare para el entrenamiento. Esto implicará cargar las imágenes en los tensores PyTorch, crear lotes de datos y realizar cualquier procesamiento previo necesario, como la normalización o el aumento de datos.

En general, el conjunto de datos "fill50k" está diseñado para usarse para entrenar modelos de aprendizaje automático para llenar círculos con colores específicos según las indicaciones proporcionadas en el conjunto de datos. Es un recurso útil para investigadores y desarrolladores que trabajan en la generación de imágenes y tareas relacionadas utilizando técnicas de aprendizaje profundo.

Cargando el conjunto de datos

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)

Tenga en cuenta que todos los pesos dentro de ControlNet también se copian desde SD para que ninguna capa se entrene desde cero y aún esté ajustando todo el modelo.

Entrenando al modelo

Este script de Python usa PyTorch Lightning para entrenar un modelo de aprendizaje automático personalizado en un conjunto de datos personalizado. Carga el modelo preentrenado, establece las configuraciones, como el tamaño del lote, la frecuencia del registrador, la tasa de aprendizaje y más. Luego crea un DataLoader para cargar datos del conjunto de datos personalizado, crea un registrador de imágenes para registrar el progreso del entrenamiento y configura PyTorch Lightning Trainer para manejar el proceso de entrenamiento. Finalmente, entrena el modelo usando Trainer y DataLoader, y registra el progreso usando 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)

      
                

                    
                    
                
            

      
                

                    
                    
                
            

      
                

                    
                    
                
            

En general, el proceso de generar una imagen a partir de una señal de aviso y control implica entrenar un modelo de aprendizaje profundo en un conjunto de datos de pares de entrada-salida y ajustar los parámetros del modelo para que aprenda a combinar las entradas para generar las salidas.

aprendizajes

Preparación de datos y modelos:

• Los modelos previamente entrenados pueden ser útiles para una variedad de tareas y se pueden integrar fácilmente en proyectos.
• Hay diferentes modelos preentrenados disponibles para controlar diferentes aspectos de una imagen, como los bordes, la profundidad y las poses humanas.
• Las bibliotecas de código abierto como PyTorch se pueden usar para entrenar e implementar modelos de aprendizaje automático.

• La difusión estable es una técnica poderosa para generar imágenes de alta calidad basadas en entradas de control.
• El uso de un modelo de difusión con una red de control puede permitir un control más preciso sobre la imagen generada.
• Se pueden usar varias indicaciones para generar diferentes imágenes basadas en la misma entrada de control.

• Las técnicas de procesamiento de imágenes como la detección de bordes Canny y OpenPose se pueden usar para extraer información específica de las imágenes, que luego se pueden usar como entradas de control para la generación de imágenes.
• Diferentes tipos de técnicas de procesamiento de imágenes son adecuadas para diferentes tipos de entradas de control, como bordes, poses y contornos.

Fuentes

Utilizamos los siguientes recursos en nuestro análisis y base de código como referencia.