Adapter un modèle à une nouvelle tâche

De nombreux systèmes de diffusion partagent les mêmes composants, ce qui vous permet d’adapter un modèle préentraîné pour une tâche à une tâche entièrement différente.

Ce guide vous montre comment adapter un modèle de texte vers image pré-entraîné pour l'inpainting en initialisant et en modifiant l'architecture d'un UNet2DConditionModel pré-entraîné.

Configurer les paramètres de UNet2DConditionModel

Par défaut, un UNet2DConditionModel accepte 4 canaux en entrée. Par exemple, chargez un modèle texte-vers-image préentraîné comme stable-diffusion-v1-5/stable-diffusion-v1-5 et vérifiez le nombre 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

L'inpainting nécessite 9 canaux dans l'échantillon d'entrée. Vous pouvez vérifier cette valeur dans un modèle d'inpainting préentraîné comme 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

Pour adapter votre modèle de texte vers image à l'inpainting, vous devez modifier le nombre de in_channels de 4 à 9.

Initialisez un UNet2DConditionModel avec les poids du modèle texte-vers-image pré-entraîné, et modifiez in_channels en 9. Modifier le nombre de in_channels implique de définir ignore_mismatched_sizes=True et low_cpu_mem_usage=False pour éviter une erreur de correspondance des tailles, car la forme est désormais différente.

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

Les poids préentraînés des autres composants du modèle texte-vers-image sont initialisés à partir de leurs points de contrôle, mais les poids des canaux d'entrée (conv_in.weight) du unet sont initialisés aléatoirement. Il est important d'ajuster finement le modèle pour l'inpainting, car sinon le modèle renvoie du bruit.