Adaptar um modelo a uma nova tarefa

Muitos sistemas de difusão compartilham os mesmos componentes, permitindo que você adapte um modelo pré-treinado para uma tarefa a uma tarefa completamente diferente.

Este guia mostrará como adaptar um modelo pré-treinado de texto para imagem para inpainting, inicializando e modificando a arquitetura de um UNet2DConditionModel pré-treinado.

Configure os parâmetros do UNet2DConditionModel

Por padrão, um UNet2DConditionModel aceita 4 canais na amostra de entrada. Por exemplo, carregue um modelo pré-treinado de texto para imagem, como stable-diffusion-v1-5/stable-diffusion-v1-5, e verifique o número de in_channels:

from diffusers import StableDiffusionPipeline

pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", use_safetensors=True)
pipeline.unet.config["in_channels"]
4

O inpainting requer 9 canais na amostra de entrada. Você pode verificar esse valor em um modelo de inpainting pré-treinado, como stable-diffusion-v1-5/stable-diffusion-inpainting:

from diffusers import StableDiffusionPipeline

pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-inpainting", use_safetensors=True)
pipeline.unet.config["in_channels"]
9

Para adaptar seu modelo de texto para imagem para inpainting, você precisará alterar o número de in_channels de 4 para 9.

Inicialize um UNet2DConditionModel com os pesos do modelo pré-treinado de texto para imagem e altere in_channels para 9. Alterar o número de in_channels significa que você precisa definir ignore_mismatched_sizes=True e low_cpu_mem_usage=False para evitar um erro de incompatibilidade de tamanho, pois a forma agora é diferente.

from diffusers import AutoModel

model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
unet = AutoModel.from_pretrained(
    model_id,
    subfolder="unet",
    in_channels=9,
    low_cpu_mem_usage=False,
    ignore_mismatched_sizes=True,
    use_safetensors=True,
)

Os pesos pré-treinados dos outros componentes do modelo de texto para imagem são inicializados a partir de seus checkpoints, mas os pesos do canal de entrada (conv_in.weight) do unet são inicializados aleatoriamente. É importante fazer o finetuning do modelo para inpainting, pois, caso contrário, o modelo retorna ruído.