Адаптация модели под новую задачу

Многие диффузионные системы используют одни и те же компоненты, что позволяет адаптировать предобученную модель для одной задачи к совершенно другой.

This guide will show you how to adapt a pretrained text-to-image model for inpainting by initializing and modifying the architecture of a pretrained UNet2DConditionModel.

Настройка параметров UNet2DConditionModel

По умолчанию UNet2DConditionModel принимает 4 канала во входном образце. Например, загрузите предобученную модель text-to-image, такую как stable-diffusion-v1-5/stable-diffusion-v1-5, и посмотрите на количество 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

Для инпейнтинга требуется 9 каналов во входном образце. Вы можете проверить это значение в предобученной модели инпейнтинга, например 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

Чтобы адаптировать вашу модель text-to-image для inpainting, вам нужно изменить количество in_channels с 4 на 9.

Инициализируйте UNet2DConditionModel с весами предобученной модели text-to-image и измените in_channels на 9. Изменение количества in_channels требует установки ignore_mismatched_sizes=True и low_cpu_mem_usage=False, чтобы избежать ошибки несовпадения размеров, так как форма теперь отличается.

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

Предобученные веса остальных компонентов модели text-to-image инициализируются из их контрольных точек, однако веса входных каналов (conv_in.weight) в unet инициализируются случайным образом. Важно дообучить модель для inpainting, иначе она будет возвращать шум.