Adaptar un modelo a una nueva tarea

Muchos sistemas de difusión comparten los mismos componentes, lo que permite adaptar un modelo preentrenado para una tarea a una tarea completamente diferente.

Esta guía le mostrará cómo adaptar un modelo de texto a imagen preentrenado para inpainting, inicializando y modificando la arquitectura de un UNet2DConditionModel preentrenado.

Configurar los parámetros de UNet2DConditionModel

Un UNet2DConditionModel acepta 4 canales en la muestra de entrada de forma predeterminada. Por ejemplo, cargue un modelo de texto a imagen preentrenado como stable-diffusion-v1-5/stable-diffusion-v1-5 y observe el 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

El inpainting requiere 9 canales en la muestra de entrada. Puedes comprobar este valor en un modelo de inpainting preentrenado 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 tu modelo de texto a imagen para inpainting, debes cambiar el número de in_channels de 4 a 9.

Inicializa un UNet2DConditionModel con los pesos del modelo de texto a imagen preentrenado y cambia in_channels a 9. Al cambiar el número de in_channels, es necesario establecer ignore_mismatched_sizes=True y low_cpu_mem_usage=False para evitar un error de incompatibilidad de tamaño, ya que la forma es diferente ahora.

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

Los pesos preentrenados de los demás componentes del modelo de texto a imagen se inicializan desde sus puntos de control, pero los pesos del canal de entrada (conv_in.weight) del unet se inicializan aleatoriamente. Es importante ajustar el modelo para inpainting, ya que de lo contrario el modelo devuelve ruido.