Adaptasi model untuk tugas baru

Banyak sistem difusi berbagi komponen yang sama, memungkinkan Anda untuk menyesuaikan model pralatih untuk satu tugas menjadi tugas yang sepenuhnya berbeda.

Panduan ini akan menunjukkan cara mengadaptasi model teks-ke-gambar yang telah dilatih sebelumnya untuk inpainting dengan menginisialisasi dan memodifikasi arsitektur dari UNet2DConditionModel yang telah dilatih sebelumnya.

Konfigurasi parameter UNet2DConditionModel

UNet2DConditionModel secara default menerima 4 channel pada input sample. Sebagai contoh, muat model text-to-image pretrained seperti stable-diffusion-v1-5/stable-diffusion-v1-5 dan perhatikan jumlah 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 membutuhkan 9 channel pada sampel input. Anda dapat memeriksa nilai ini pada model inpainting pretrained seperti 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

Untuk menyesuaikan model teks-ke-gambar Anda untuk inpainting, Anda perlu mengubah jumlah in_channels dari 4 menjadi 9.

Inisialisasi UNet2DConditionModel dengan bobot model teks-ke-gambar yang telah dilatih sebelumnya, dan ubah in_channels menjadi 9. Mengubah jumlah in_channels berarti Anda perlu mengatur ignore_mismatched_sizes=True dan low_cpu_mem_usage=False untuk menghindari error ketidakcocokan ukuran karena bentuknya sekarang berbeda.

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

Bobot pra-terlatih dari komponen lain pada model teks-ke-gambar diinisialisasi dari checkpoint mereka, tetapi bobot kanal input (conv_in.weight) dari unet diinisialisasi secara acak. Penting untuk melakukan finetune pada model untuk inpainting karena jika tidak, model akan mengembalikan noise.