Điều chỉnh một mô hình cho một tác vụ mới

Nhiều hệ thống khuếch tán chia sẻ các thành phần giống nhau, cho phép bạn điều chỉnh một mô hình được huấn luyện trước cho một tác vụ sang một tác vụ hoàn toàn khác.

Hướng dẫn này sẽ cho bạn thấy cách điều chỉnh một mô hình text-to-image được huấn luyện trước cho tác vụ inpainting bằng cách khởi tạo và sửa đổi kiến trúc của một UNet2DConditionModel đã được huấn luyện trước.

Cấu hình các tham số của UNet2DConditionModel

UNet2DConditionModel mặc định nhận 4 kênh trong mẫu đầu vào. Ví dụ, hãy tải một mô hình text-to-image được huấn luyện trước như stable-diffusion-v1-5/stable-diffusion-v1-5 và xem số lượng 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 yêu cầu 9 kênh trong mẫu đầu vào. Bạn có thể kiểm tra giá trị này trong một mô hình inpainting được huấn luyện trước như 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

Để điều chỉnh mô hình text-to-image của bạn cho inpainting, bạn cần thay đổi số lượng in_channels từ 4 thành 9.

Khởi tạo một UNet2DConditionModel với trọng số mô hình text-to-image được huấn luyện trước, và thay đổi in_channels thành 9. Việc thay đổi số lượng in_channels có nghĩa là bạn cần đặt ignore_mismatched_sizes=True và low_cpu_mem_usage=False để tránh lỗi không khớp kích thước vì hình dạng hiện đã khác.

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

Trọng số tiền huấn luyện của các thành phần khác từ mô hình text-to-image được khởi tạo từ checkpoint của chúng, nhưng trọng số kênh đầu vào (conv_in.weight) của unet được khởi tạo ngẫu nhiên. Việc tinh chỉnh mô hình cho inpainting là rất quan trọng vì nếu không, mô hình sẽ trả về nhiễu.