#
Adapt a model to a new task

Many diffusion systems share the same components, allowing you to adapt a pretrained model for one task to an entirely different task.

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.

##
Configure UNet2DConditionModel parameters

A UNet2DConditionModel by default accepts 4 channels in the input sample. For example, load a pretrained text-to-image model like `stable-diffusion-v1-5/stable-diffusion-v1-5` and take a look at the number of `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

Inpainting requires 9 channels in the input sample. You can check this value in a pretrained inpainting model like `runwayml/stable-diffusion-inpainting`:

    from diffusers import StableDiffusionPipeline

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

To adapt your text-to-image model for inpainting, you’ll need to change the number of `in_channels` from 4 to 9.

Initialize a UNet2DConditionModel with the pretrained text-to-image model weights, and change `in_channels` to 9. Changing the number of `in_channels` means you need to set `ignore_mismatched_sizes=True` and `low_cpu_mem_usage=False` to avoid a size mismatch error because the shape is different now.

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

The pretrained weights of the other components from the text-to-image model are initialized from their checkpoints, but the input channel weights (`conv_in.weight`) of the `unet` are randomly initialized. It is important to finetune the model for inpainting because otherwise the model returns noise.

< > Update on GitHub