Entendendo pipelines, modelos e agendadores

🧨 O Diffusers foi projetado para ser uma caixa de ferramentas amigável e flexível para a construção de sistemas de difusão adaptados ao seu caso de uso. No núcleo da caixa de ferramentas estão os modelos e os agendadores. Embora o DiffusionPipeline agrupe esses componentes para facilitar o uso, você também pode desmembrar o pipeline e usar os modelos e os agendadores separadamente para criar novos sistemas de difusão.

Neste tutorial, você aprenderá a usar modelos e agendadores para montar um sistema de difusão para inferência, começando com um pipeline básico e, em seguida, avançando para o pipeline do Stable Diffusion.

Desconstruir um pipeline básico

Uma pipeline é uma maneira rápida e fácil de executar um modelo para inferência, exigindo no máximo quatro linhas de código para gerar uma imagem:

>>> from diffusers import DDPMPipeline

>>> ddpm = DDPMPipeline.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
>>> image = ddpm(num_inference_steps=25).images[0]
>>> image

Foi super fácil, mas como o pipeline fez isso? Vamos decompor o pipeline e dar uma olhada no que está acontecendo por baixo dos panos.

No exemplo acima, o pipeline contém um modelo UNet2DModel e um DDPMScheduler. O pipeline remove o ruído de uma imagem ao pegar ruído aleatório do tamanho da saída desejada e passá-lo pelo modelo várias vezes. Em cada timestep, o modelo prevê o resíduo de ruído e o scheduler usa isso para prever uma imagem com menos ruído. O pipeline repete esse processo até atingir o final do número especificado de etapas de inferência.

Para recriar o pipeline com o modelo e o agendador separadamente, vamos escrever nosso próprio processo de remoção de ruído.

1. Carregue o modelo e o agendador:

>>> from diffusers import DDPMScheduler, UNet2DModel

>>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
>>> model = UNet2DModel.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

2. Defina o número de passos de tempo para executar o processo de remoção de ruído:

>>> scheduler.set_timesteps(50)

3. Definir os timesteps do scheduler cria um tensor com elementos igualmente espaçados, 50 neste exemplo. Cada elemento corresponde a um timestep no qual o modelo remove o ruído de uma imagem. Quando você criar o loop de remoção de ruído mais tarde, iterará sobre este tensor para remover o ruído de uma imagem:

>>> scheduler.timesteps
tensor([980, 960, 940, 920, 900, 880, 860, 840, 820, 800, 780, 760, 740, 720,
    700, 680, 660, 640, 620, 600, 580, 560, 540, 520, 500, 480, 460, 440,
    420, 400, 380, 360, 340, 320, 300, 280, 260, 240, 220, 200, 180, 160,
    140, 120, 100,  80,  60,  40,  20,   0])

4. Crie um ruído aleatório com a mesma forma que a saída desejada:

>>> import torch

>>> sample_size = model.config.sample_size
>>> noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")  # or "mps", "xpu", "cpu"

5. Agora escreva um loop para iterar pelos timesteps. Em cada timestep, o modelo executa uma passagem UNet2DModel.forward() e retorna o residual ruidoso. O método step() do scheduler recebe o residual ruidoso, o timestep e a entrada, e prevê a imagem no timestep anterior. Essa saída se torna a próxima entrada do modelo no loop de remoção de ruído, e o processo se repete até atingir o final do array timesteps.

>>> input = noise

>>> for t in scheduler.timesteps:
...     with torch.no_grad():
...         noisy_residual = model(input, t).sample
...     previous_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
...     input = previous_noisy_sample

Este é todo o processo de remoção de ruído, e você pode usar esse mesmo padrão para escrever qualquer sistema de difusão.

6. A última etapa é converter a saída desruída em uma imagem:

>>> from PIL import Image
>>> import numpy as np

>>> image = (input / 2 + 0.5).clamp(0, 1).squeeze()
>>> image = (image.permute(1, 2, 0) * 255).round().to(torch.uint8).cpu().numpy()
>>> image = Image.fromarray(image)
>>> image

Na próxima seção, você vai colocar suas habilidades à prova e analisar o pipeline mais complexo do Stable Diffusion. Os passos são mais ou menos os mesmos. Você vai inicializar os componentes necessários e definir o número de timesteps para criar um array timestep. O array timestep é usado no loop de remoção de ruído, e para cada elemento deste array, o modelo prevê uma imagem com menos ruído. O loop de remoção de ruído itera pelos timesteps e, em cada timestep, ele produz um residual com ruído e o scheduler o usa para prever uma imagem com menos ruído no timestep anterior. Esse processo é repetido até você chegar ao final do array timestep.

Vamos testar!

Desconstrua o pipeline do Stable Diffusion

Stable Diffusion é um modelo de latent diffusion de texto para imagem. É chamado de modelo de latent diffusion porque trabalha com uma representação de menor dimensão da imagem, em vez do espaço de pixels real, o que o torna mais eficiente em termos de memória. O encoder comprime a imagem em uma representação menor, e um decoder converte a representação comprimida de volta em uma imagem. Para modelos de texto para imagem, você precisará de um tokenizer e de um encoder para gerar embeddings de texto. A partir do exemplo anterior, você já sabe que precisa de um modelo UNet e de um scheduler.

Como você pode ver, isso já é mais complexo que o pipeline DDPM, que contém apenas um modelo UNet. O modelo Stable Diffusion possui três modelos pré-treinados separados.

💡 Leia o blog Como o Stable Diffusion funciona? para mais detalhes sobre como funcionam os modelos VAE, UNet e text encoder.

Agora que você sabe o que precisa para o pipeline do Stable Diffusion, carregue todos esses componentes com o método from_pretrained(). Você pode encontrá-los no checkpoint pré-treinado stable-diffusion-v1-5/stable-diffusion-v1-5, e cada componente é armazenado em uma subpasta separada:

>>> from PIL import Image
>>> import torch
>>> from transformers import CLIPTextModel, CLIPTokenizer
>>> from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler

>>> vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", use_safetensors=True)
>>> tokenizer = CLIPTokenizer.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="tokenizer")
>>> text_encoder = CLIPTextModel.from_pretrained(
...     "CompVis/stable-diffusion-v1-4", subfolder="text_encoder", use_safetensors=True
... )
>>> unet = UNet2DConditionModel.from_pretrained(
...     "CompVis/stable-diffusion-v1-4", subfolder="unet", use_safetensors=True
... )

Em vez do PNDMScheduler padrão, substitua-o pelo UniPCMultistepScheduler para ver como é fácil plugar um scheduler diferente:

>>> from diffusers import UniPCMultistepScheduler

>>> scheduler = UniPCMultistepScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")

Para acelerar a inferência, mova os modelos para uma GPU, pois, ao contrário do agendador, eles possuem pesos treináveis:

>>> torch_device = "cuda"  # or "mps", "xpu", "cpu"
>>> vae.to(torch_device)
>>> text_encoder.to(torch_device)
>>> unet.to(torch_device)

Criar embeddings de texto

O próximo passo é tokenizar o texto para gerar embeddings. O texto é usado para condicionar o modelo UNet e direcionar o processo de difusão para algo que se assemelhe ao prompt de entrada.

💡 O parâmetro guidance_scale determina o quanto de peso deve ser dado ao prompt ao gerar uma imagem.

Sinta-se à vontade para escolher qualquer prompt que quiser, caso queira gerar algo diferente!

>>> prompt = ["a photograph of an astronaut riding a horse"]
>>> height = 512  # default height of Stable Diffusion
>>> width = 512  # default width of Stable Diffusion
>>> num_inference_steps = 25  # Number of denoising steps
>>> guidance_scale = 7.5  # Scale for classifier-free guidance
>>> generator = torch.manual_seed(0)  # Seed generator to create the initial latent noise
>>> batch_size = len(prompt)

Tokenize o texto e gere os embeddings a partir do prompt:

>>> 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]

Você também precisará gerar os embeddings de texto incondicionais, que são os embeddings para o token de padding. Eles precisam ter a mesma forma (batch_size e seq_length) que os text_embeddings condicionais:

>>> max_length = text_input.input_ids.shape[-1]
>>> uncond_input = tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt")
>>> uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]

Vamos concatenar os embeddings condicionais e incondicionais em um lote para evitar duas passadas de frente:

>>> text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

Criar ruído aleatório

Em seguida, gere um ruído aleatório inicial como ponto de partida para o processo de difusão. Essa é a representação latente da imagem, que será gradualmente desruidada. Neste ponto, a imagem latent é menor que o tamanho final da imagem, mas isso não é um problema, pois o modelo a transformará nas dimensões finais de 512x512 mais tarde.

💡 A altura e a largura são divididas por 8 porque o modelo vae possui 3 camadas de redução de amostragem. Você pode verificar executando o seguinte:

Copiado

2 ** (len(vae.config.block_out_channels) - 1) == 8

>>> latents = torch.randn(
...     (batch_size, unet.config.in_channels, height // 8, width // 8),
...     generator=generator,
...     device=torch_device,
... )

Remova o ruído da imagem

Comece escalando a entrada com a distribuição inicial de ruído, sigma, o valor de escala do ruído, que é necessário para agendadores aprimorados como UniPCMultistepScheduler:

>>> latents = latents * scheduler.init_noise_sigma

A última etapa é criar o loop de remoção de ruído que transformará progressivamente o ruído puro em latents em uma imagem descrita pelo seu prompt. Lembre-se, o loop de remoção de ruído precisa fazer três coisas:

1. Defina os timesteps do scheduler a serem usados durante o desruidamento.
2. Itere sobre os timesteps.
3. Em cada passo de tempo, chame o modelo UNet para prever o resíduo de ruído e passe-o ao agendador para calcular a amostra ruidosa anterior.

>>> from tqdm.auto import tqdm

>>> scheduler.set_timesteps(num_inference_steps)

>>> for t in tqdm(scheduler.timesteps):
...     # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
...     latent_model_input = torch.cat([latents] * 2)

...     latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)

...     # predict the noise residual
...     with torch.no_grad():
...         noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

...     # perform guidance
...     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
...     noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

...     # compute the previous noisy sample x_t -> x_t-1
...     latents = scheduler.step(noise_pred, t, latents).prev_sample

Decodifique a imagem

A etapa final é usar o vae para decodificar a representação latente em uma imagem e obter a saída decodificada com sample:

# scale and decode the image latents with vae
latents = 1 / 0.18215 * latents
with torch.no_grad():
    image = vae.decode(latents).sample

Por fim, converta a imagem para PIL.Image para ver a imagem gerada!

>>> image = (image / 2 + 0.5).clamp(0, 1).squeeze()
>>> image = (image.permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()
>>> image = Image.fromarray(image)
>>> image

Próximas etapas

De pipelines básicos a complexos, você viu que tudo o que realmente precisa para escrever seu próprio sistema de difusão é um loop de desruidamento. O loop deve definir os timesteps do scheduler, iterar sobre eles e alternar entre chamar o modelo UNet para prever o resíduo de ruído e passá-lo ao scheduler para calcular a amostra ruidosa anterior.

É exatamente para isso que o 🧨 Diffusers foi projetado: para tornar intuitivo e fácil escrever o seu próprio sistema de difusão usando modelos e agendadores.

Para os próximos passos, sinta-se à vontade para:

• Aprenda como construir e contribuir com um pipeline no 🧨 Diffusers. Mal podemos esperar para ver o que você vai criar!
• Explore os pipelines existentes na biblioteca e veja se você consegue deconstruí-los e construir um pipeline do zero usando os modelos e os agendadores separadamente.