Controle os modelos de difusão adicionando condições extras

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

Os modelos de difusão de texto para imagem fizeram progressos impressionantes nos últimos anos na geração de imagens de alta qualidade a partir de descrições textuais. Esses modelos aprendem a transformar a linguagem natural em espaço de pixel, permitindo gerar imagens realistas que correspondem à descrição de entrada. No entanto, ainda há espaço para melhorias em termos de controle sobre as imagens geradas. Adicionar controle condicional a modelos de difusão de texto para imagem pode permitir que os usuários orientem a geração de imagens com mais precisão, resultando em resultados mais precisos e personalizados. Nesse contexto, o uso de esboços, contornos e poses como meio de controle condicional está se tornando cada vez mais popular. Neste blog, exploraremos o conceito de controle condicional em modelos de difusão de texto para imagem,

Neste blog, exploraremos o conceito de combinar ControlNet com difusão estável e seu potencial para aprimorar a geração de texto para imagem.

O que é difusão estável?

A difusão estável é uma variante do modelo de difusão, um modelo generativo que aprende a gerar imagens de alta qualidade modelando a distribuição de fontes intermediárias de ruído. O modelo de difusão é um processo sequencial que começa com uma distribuição fixa de ruído e aplica uma sequência de transformações, cada uma das quais modifica a distribuição de ruído para produzir uma nova distribuição que corresponda melhor à distribuição alvo. A distribuição final é usada para gerar a imagem de saída. O modelo de difusão tem se mostrado eficaz na geração de imagens de alta qualidade, mas pode ser difícil de treinar devido à necessidade de estimar distribuições condicionais complexas.

A difusão é o processo que ocorre dentro do componente rosa “criador de informações da imagem”. Tendo as incorporações de token que representam o texto de entrada e uma matriz de informações de imagem inicial aleatória (também chamadas latentes ), o processo produz uma matriz de informações que o decodificador de imagem usa para pintar a imagem final.

Este procedimento é realizado em etapas. Cada etapa adiciona dados mais relevantes. Podemos inspecionar a matriz de latentes aleatórios e ver como ela se traduz em ruído visual para ter uma noção do processo. Nesse caso, a inspeção visual implica em executá-lo por meio de um decodificador de imagem. A difusão ocorre em várias etapas, com cada etapa operando em uma matriz de latentes de entrada e produzindo outra matriz de latentes que se assemelha mais ao texto de entrada, bem como a todas as informações visuais que o modelo captou de todas as imagens nas quais o modelo foi treinado. Um conjunto desses latentes pode ser visualizado para ver quais informações são adicionadas em cada etapa.

U sing difusão estável em python usando o pipeline 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 é uma estrutura de rede neural que permite o controle explícito sobre grandes modelos de difusão pré-treinados, permitindo que condições de entrada adicionais sejam suportadas. O ControlNet foi projetado para aprender condições específicas de tarefas de maneira ponta a ponta, tornando o processo de aprendizado robusto mesmo quando o conjunto de dados de treinamento é pequeno, com menos de 50 mil amostras. Esse recurso torna o ControlNet uma ferramenta poderosa para geração de texto para imagem, onde as variáveis ​​de controle podem ser usadas para especificar os atributos desejados das imagens geradas, como forma, cor ou textura.

Como funciona

O ControlNet funciona copiando os pesos dos blocos da rede neural em duas cópias: uma cópia “bloqueada” e uma cópia “treinável”. A cópia “bloqueada” preserva o modelo original, enquanto a cópia “treinável” é usada para aprender as condições especificadas.

Durante o treinamento, a ControlNet usa uma técnica chamada “convolução zero”, que é uma convolução 1x1 com peso e bias inicializados como zeros. Antes do treinamento, todas as convoluções de zero produzem zeros e o ControlNet não causa nenhuma distorção no modelo original. A cópia “treinável” aprende as condições especificadas, enquanto a cópia “bloqueada” permanece inalterada.

Esta abordagem tem várias vantagens. Primeiro, permite o treinamento em pequenos conjuntos de dados de pares de imagens sem comprometer os modelos de difusão prontos para produção. Em segundo lugar, nenhuma camada é treinada do zero, o que significa que o modelo original é seguro. Em terceiro lugar, esta abordagem permite treinamento em pequena escala ou mesmo em dispositivos pessoais, o que a torna prática para uma ampla gama de aplicações.

Usando o pipeline StableDiffusion ControlNet

estamos usando um Pipeline para geração de texto para imagem usando Stable Diffusion com orientação ControlNet

Modelos pré-treinados disponíveis:

• Controle usando detecção de borda : lllyasviel/sd-controlnet-canny: treinado com detecção de borda inteligente, usado para controlar bordas de imagem.
• Controle usando detecção de pose: lllyasviel/sd-controlnet_openpose: treinado com a imagem óssea OpenPose, usada para controlar poses humanas.
• Controle usando Detecção de Rabisco: lllyasviel/sd-controlnet_scribble: Treinado com rabiscos humanos, usado para controlar contornos de imagem.
• lllyasviel/sd-controlnet-depth: treinado com estimativa de profundidade Midas, usado para controlar a profundidade da imagem.
• lllyasviel/sd-controlnet-hed: treinado com detecção de borda HED (borda suave), usado para controlar as bordas suaves da imagem.
• lllyasviel/sd-controlnet-mlsd: treinado com detecção de linha M-LSD, usado para controlar linhas retas em uma imagem.
• lllyasviel/sd-controlnet-normal: Treinado com mapa normal, usado para controlar imagens normais.
• lllyasviel/sd-controlnet_seg: Treinado com segmentação semântica, usado para controlar a segmentação semântica da imagem.

Primeiro, temos que configurar nosso ambiente e importar a difusão estável do rosto de abraço, aqui estamos usando a difusão estável-v1–5

O código instala alguns pacotes usando pip e depois importa módulos e classes desses pacotes.

A primeira linha instala os seguintes pacotes:

• diffusers==0.14.0: Um pacote para implementar modelos de difusão no PyTorch.
• transformers: um pacote para tarefas de processamento de linguagem natural (NLP), como classificação de texto e geração de linguagem.
• xformers: um pacote para modelos de transformadores, que são um tipo de arquitetura de rede neural usada em tarefas de NLP.
• git+https://github.com/huggingface/accelerate.git: Um pacote para otimizar o treinamento de modelos de aprendizado 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

      
                

                    
                    
                
            

      
                

                    
                    
                
            

A saída do pipeobjeto é um conjunto de quatro imagens geradas com base nos prompts e nos parâmetros de controle fornecidos. Os prompts são fornecidos como uma lista de quatro strings, sendo cada string um nome de celebridade seguido por uma mensagem indicando as qualidades de imagem desejadas. As mensagens de prompt indicam um desejo por imagens de "melhor qualidade" e "extremamente detalhadas". O negative_promptparâmetro é definido para indicar as qualidades de imagem desejadas de forma negativa, solicitando imagens "monocromáticas", "lowres", "anatomicidade ruim", "pior qualidade" e "baixa qualidade". O generatorparâmetro é usado para controlar a aleatoriedade das imagens geradas e o num_inference_stepsparâmetro define o número de iterações usadas para gerar cada imagem.

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)

      
                

                    
                    
                
            

Esse código carrega um OpenposeDetectormodelo pré-treinado do hub de modelo Hugging Face usando o from_pretrained()método e o aplica a uma imagem para detectar poses humanas. Ele também carrega um pré-treinado ControlNetModele um StableDiffusionControlNetPipelinedo diffuserspacote, define alguns parâmetros de controle para o pipeline e permite algumas otimizações de memória e desempenho.

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 como podemos treinar um modelo ControlNet em um conjunto de dados. Então, para atingir nosso objetivo de controlar o SD, assumimos a tarefa de preencher as cores nos círculos. Nossa abordagem será usar uma tarefa simples para demonstrar como o ControlNet pode aprender as condições específicas da tarefa de uma maneira completa. Usaremos um prompt que descreve nosso alvo e inclui uma “Imagem de controle (imagem de origem)”.

A difusão estável já foi treinada em bilhões de imagens e já sabe o que são cores e formas. No entanto, ele não entende o significado de “Control Image (Source Image)” em nosso prompt. Ao adicionar o ControlNet, podemos ensinar o SD a entender essa condição de entrada e preencher o círculo com a cor especificada.

O conjunto de dados

O conjunto de dados “fill50k” contém uma coleção de 50.000 pares de imagens de círculos, que são armazenados em duas pastas: “source” e “target”. A pasta “source” contém imagens de círculos apenas com linhas, enquanto a pasta “target” contém imagens de círculos preenchidos com uma cor específica.

Além das imagens, o conjunto de dados inclui um arquivo “prompt.json”, que fornece informações sobre cada imagem. Cada prompt tem a forma de uma descrição do círculo e a cor do plano de fundo. Por exemplo, um prompt pode ler “um círculo azul em um fundo amarelo”.

Para usar este conjunto de dados com o PyTorch, você precisará escrever um script que leia as imagens e seus prompts correspondentes da pasta “fill50k” e os prepare para o treinamento. Isso envolverá carregar as imagens em tensores PyTorch, criar lotes de dados e realizar qualquer pré-processamento necessário, como normalização ou aumento de dados.

No geral, o conjunto de dados “fill50k” foi projetado para ser usado para treinar modelos de aprendizado de máquina para preencher círculos com cores específicas com base nos prompts fornecidos no conjunto de dados. É um recurso útil para pesquisadores e desenvolvedores que trabalham na geração de imagens e tarefas relacionadas usando técnicas de aprendizado profundo.

Carregando o conjunto de dados

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)

Observe que todos os pesos dentro do ControlNet também são copiados do SD para que nenhuma camada seja treinada do zero e você ainda esteja ajustando todo o modelo.

Treinando o modelo

Este script Python usa o PyTorch Lightning para treinar um modelo de aprendizado de máquina personalizado em um conjunto de dados personalizado. Ele carrega o modelo pré-treinado, define as configurações como tamanho do lote, frequência do registrador, taxa de aprendizado e muito mais. Em seguida, ele cria um DataLoader para carregar dados do conjunto de dados personalizado, cria um registrador de imagem para registrar o progresso do treinamento e configura o PyTorch Lightning Trainer para lidar com o processo de treinamento. Por fim, ele treina o modelo usando o Trainer e o DataLoader e registra o progresso usando o 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)

      
                

                    
                    
                
            

      
                

                    
                    
                
            

      
                

                    
                    
                
            

No geral, o processo de geração de uma imagem a partir de um sinal de prompt e controle envolve o treinamento de um modelo de aprendizado profundo em um conjunto de dados de pares de entrada-saída e o ajuste dos parâmetros do modelo para que ele aprenda a combinar as entradas para gerar as saídas.

Aprendizados

Preparação de dados e modelo:

• Modelos pré-treinados podem ser úteis para uma variedade de tarefas e podem ser facilmente integrados a projetos.
• Diferentes modelos pré-treinados estão disponíveis para controlar diferentes aspectos de uma imagem, como bordas, profundidade e poses humanas.
• Bibliotecas de código aberto como o PyTorch podem ser usadas para treinamento e implantação de modelos de aprendizado de máquina.

• A difusão estável é uma técnica poderosa para gerar imagens de alta qualidade com base em entradas de controle.
• Usar um modelo de difusão com uma rede de controle pode permitir um controle mais preciso sobre a imagem gerada.
• Vários prompts podem ser usados ​​para gerar imagens diferentes com base na mesma entrada de controle.

• Técnicas de processamento de imagem como detecção de borda Canny e OpenPose podem ser usadas para extrair informações específicas de imagens, que podem então ser usadas como entradas de controle para geração de imagens.
• Diferentes tipos de técnicas de processamento de imagem são adequados para diferentes tipos de entradas de controle, como bordas, poses e contornos.

Fontes

Utilizamos os seguintes recursos em nossa análise e base de código para referência.