Hugging Face's logo

Diffusers documentation

Custom Diffusion

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Custom Diffusion

Custom Diffusion là một kỹ thuật huấn luyện để cá nhân hóa các mô hình tạo ảnh. Giống như Textual Inversion, DreamBooth và LoRA, Custom Diffusion chỉ cần một vài (~4-5) hình ảnh ví dụ. Kỹ thuật này hoạt động bằng cách chỉ huấn luyện trọng số trong các lớp cross-attention, và nó sử dụng một từ đặc biệt để đại diện cho khái niệm mới được học. Custom Diffusion độc đáo ở chỗ nó có thể học nhiều khái niệm cùng một lúc.

Nếu bạn đang huấn luyện trên GPU có vRAM hạn chế, bạn nên thử bật xFormers bằng --enable_xformers_memory_efficient_attention để huấn luyện nhanh hơn với yêu cầu vRAM thấp hơn (16GB). Để tiết kiệm thêm bộ nhớ, hãy thêm --set_grads_to_none vào đối số huấn luyện để đặt gradient thành None thay vì số 0 (tùy chọn này có thể gây ra một số vấn đề, vì vậy nếu bạn gặp sự cố, hãy thử xóa tham số này).

Hướng dẫn này sẽ khám phá script train_custom_diffusion.py để giúp bạn làm quen hơn với nó, cũng như cách bạn có thể điều chỉnh nó cho trường hợp sử dụng của riêng mình.

Trước khi chạy script, hãy đảm bảo bạn cài đặt thư viện từ mã nguồn:

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

Điều hướng đến thư mục chứa ví dụ với script huấn luyện và cài đặt các phụ thuộc cần thiết:

cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval

🤗 Accelerate là một thư viện giúp bạn huấn luyện trên nhiều GPU/TPU hoặc với độ chính xác hỗn hợp. Nó sẽ tự động cấu hình thiết lập huấn luyện của bạn dựa trên phần cứng và môi trường của bạn. Hãy xem 🤗 Accelerate Quick tour để tìm hiểu thêm.

Khởi tạo một môi trường 🤗 Accelerate:

accelerate config

Để thiết lập môi trường 🤗 Accelerate mặc định mà không cần chọn bất kỳ cấu hình nào:

accelerate config default

Hoặc nếu môi trường của bạn không hỗ trợ shell tương tác, chẳng hạn như notebook, bạn có thể sử dụng:

from accelerate.utils import write_basic_config

write_basic_config()

Cuối cùng, nếu bạn muốn huấn luyện một mô hình trên bộ dữ liệu của riêng mình, hãy xem hướng dẫn Tạo bộ dữ liệu để huấn luyện để tìm hiểu cách tạo một bộ dữ liệu tương thích với script huấn luyện.

Các phần tiếp theo làm nổi bật những phần quan trọng trong script huấn luyện để hiểu cách chỉnh sửa nó, nhưng không đi sâu vào mọi khía cạnh của script. Nếu bạn muốn tìm hiểu thêm, hãy tự do đọc qua script và cho chúng tôi biết nếu bạn có bất kỳ câu hỏi hoặc lo ngại nào.

Tham số của script

Tệp script huấn luyện chứa tất cả các tham số giúp bạn tùy chỉnh phiên huấn luyện của mình. Chúng được tìm thấy trong hàm parse_args(). Hàm này đi kèm với các giá trị mặc định, nhưng bạn cũng có thể tự đặt giá trị của riêng mình trong lệnh huấn luyện nếu muốn.

Ví dụ, để thay đổi độ phân giải của hình ảnh đầu vào:

accelerate launch train_custom_diffusion.py \
  --resolution=256

Nhiều tham số cơ bản được mô tả trong hướng dẫn huấn luyện DreamBooth, vì vậy hướng dẫn này tập trung vào các tham số đặc thù của Custom Diffusion:

  • --freeze_model: đóng băng các tham số key và value trong lớp cross-attention; mặc định là crossattn_kv, nhưng bạn có thể đặt thành crossattn để huấn luyện tất cả các tham số trong lớp cross-attention
  • --concepts_list: để học nhiều khái niệm, hãy cung cấp đường dẫn đến tệp JSON chứa các khái niệm
  • --modifier_token: một từ đặc biệt được sử dụng để đại diện cho khái niệm đã học
  • --initializer_token: một từ đặc biệt được sử dụng để khởi tạo các embedding của modifier_token

Mất mát bảo toàn tiên nghiệm

Mất tích lũy tiên nghiệm là một phương pháp sử dụng các mẫu do chính mô hình tạo ra để hỗ trợ việc học cách tạo ra các hình ảnh đa dạng hơn. Vì các hình ảnh mẫu được tạo ra này thuộc cùng một lớp với các hình ảnh bạn đã cung cấp, chúng giúp mô hình giữ lại những gì đã học được về lớp đó và cách sử dụng kiến thức hiện có về lớp để tạo ra các bố cục mới.

Nhiều tham số cho tổn thất bảo toàn trước được mô tả trong hướng dẫn huấn luyện DreamBooth.

Chống quá khớp

Custom Diffusion bao gồm việc huấn luyện các hình ảnh đích với một bộ nhỏ các hình ảnh thực để ngăn chặn hiện tượng quá khớp. Như bạn có thể hình dung, điều này có thể dễ dàng thực hiện khi bạn chỉ huấn luyện trên một vài hình ảnh! Tải xuống 200 hình ảnh thực bằng clip_retrieval. class_prompt nên cùng danh mục với các hình ảnh đích. Những hình ảnh này được lưu trữ trong class_data_dir.

python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200

Để bật điều hòa, hãy thêm các tham số sau:

  • --with_prior_preservation: có sử dụng prior preservation loss hay không
  • --prior_loss_weight: kiểm soát mức độ ảnh hưởng của tổn thất bảo toàn tiên nghiệm lên mô hình
  • --real_prior: có sử dụng một tập nhỏ các ảnh thực để ngăn chặn quá khớp hay không
accelerate launch train_custom_diffusion.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="./real_reg/samples_cat" \
  --class_prompt="cat" \
  --real_prior=True \

Script huấn luyện

Nhiều phần mã trong script huấn luyện Custom Diffusion tương tự như script DreamBooth. Hướng dẫn này thay vào đó tập trung vào các phần mã liên quan đến Custom Diffusion.

Tải bản dịch

  • CustomDiffusionDataset: tiền xử lý các hình ảnh, hình ảnh lớp và prompt cho việc huấn luyện
  • PromptDataset: chuẩn bị các prompt để tạo ảnh cho các lớp

Tiếp theo, modifier_token được thêm vào tokenizer, chuyển đổi thành các token id, và các token embeddings được điều chỉnh kích thước để tính đến modifier_token mới. Sau đó, embeddings của modifier_token được khởi tạo bằng embeddings của initializer_token. Tất cả các tham số trong text encoder đều được đóng băng, ngoại trừ token embeddings vì đây là thứ mà mô hình đang cố gắng học để liên kết với các khái niệm.

params_to_freeze = itertools.chain(
    text_encoder.text_model.encoder.parameters(),
    text_encoder.text_model.final_layer_norm.parameters(),
    text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)

Bây giờ bạn cần thêm Custom Diffusion weights vào các lớp attention. Đây là một bước rất quan trọng để đảm bảo hình dạng và kích thước của các trọng số attention chính xác, đồng thời thiết lập số lượng attention processors phù hợp trong mỗi khối UNet.

st = unet.state_dict()
for name, _ in unet.attn_processors.items():
    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
    if name.startswith("mid_block"):
        hidden_size = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"):
        block_id = int(name[len("up_blocks.")])
        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
    elif name.startswith("down_blocks"):
        block_id = int(name[len("down_blocks.")])
        hidden_size = unet.config.block_out_channels[block_id]
    layer_name = name.split(".processor")[0]
    weights = {
        "to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
        "to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
    }
    if train_q_out:
        weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
        weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
        weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
    if cross_attention_dim is not None:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=train_kv,
            train_q_out=train_q_out,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        ).to(unet.device)
        custom_diffusion_attn_procs[name].load_state_dict(weights)
    else:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=False,
            train_q_out=False,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        )
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)

optimizer được khởi tạo để cập nhật các tham số của lớp cross-attention:

optimizer = optimizer_class(
    itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
    if args.modifier_token is not None
    else custom_diffusion_layers.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Trong vòng lặp huấn luyện, điều quan trọng là chỉ cập nhật các embedding cho khái niệm bạn đang cố gắng học. Điều này có nghĩa là đặt gradient của tất cả các token embedding khác bằng không:

if args.modifier_token is not None:
    if accelerator.num_processes > 1:
        grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
    else:
        grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
    index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
    for i in range(len(modifier_token_id[1:])):
        index_grads_to_zero = index_grads_to_zero & (
            torch.arange(len(tokenizer)) != modifier_token_id[i]
        )
    grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
        index_grads_to_zero, :
    ].fill_(0)

Chạy script

Khi bạn đã thực hiện tất cả các thay đổi hoặc hài lòng với cấu hình mặc định, bạn đã sẵn sàng để chạy script huấn luyện! 🚀

Trong hướng dẫn này, bạn sẽ tải xuống và sử dụng các ảnh mèo mẫu này. Bạn cũng có thể tự tạo và sử dụng bộ dữ liệu của riêng mình nếu muốn (xem hướng dẫn Tạo bộ dữ liệu để huấn luyện).

Đặt biến môi trường MODEL_NAME thành id của một model trên Hub hoặc đường dẫn đến một model cục bộ, INSTANCE_DIR thành đường dẫn nơi bạn vừa tải xuống các hình ảnh mèo, và OUTPUT_DIR thành nơi bạn muốn lưu model. Bạn sẽ sử dụng <new1> làm từ đặc biệt để gắn các embeddings mới học được. Script sẽ tạo và lưu các checkpoint của model cùng tệp pytorch_custom_diffusion_weights.bin vào kho của bạn.

Để theo dõi tiến độ huấn luyện với Weights and Biases, hãy thêm tham số --report_to=wandb vào lệnh huấn luyện và chỉ định một prompt xác thực bằng --validation_prompt. Điều này hữu ích cho việc gỡ lỗi và lưu trữ các kết quả trung gian.

Nếu bạn đang huấn luyện trên khuôn mặt người, nhóm Custom Diffusion đã phát hiện ra rằng các tham số sau hoạt động tốt:

  • --learning_rate=5e-6
  • --max_train_steps có thể nằm trong khoảng từ 1000 đến 2000
  • --freeze_model=crossattn
  • sử dụng ít nhất 15-20 hình ảnh để huấn luyện
một khái niệm duy nhất
nhiều khái niệm
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"
export INSTANCE_DIR="./data/cat"

accelerate launch train_custom_diffusion.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --class_data_dir=./real_reg/samples_cat/ \
  --with_prior_preservation \
  --real_prior \
  --prior_loss_weight=1.0 \
  --class_prompt="cat" \
  --num_class_images=200 \
  --instance_prompt="photo of a <new1> cat"  \
  --resolution=512  \
  --train_batch_size=2  \
  --learning_rate=1e-5  \
  --lr_warmup_steps=0 \
  --max_train_steps=250 \
  --scale_lr \
  --hflip  \
  --modifier_token "<new1>" \
  --validation_prompt="<new1> cat sitting in a bucket" \
  --report_to="wandb" \
  --push_to_hub

Khi quá trình huấn luyện hoàn tất, bạn có thể sử dụng mô hình Custom Diffusion mới của mình để suy luận.

một khái niệm duy nhất
nhiều khái niệm
import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion("path-to-save-model", weight_name="<new1>.bin")

image = pipeline(
    "<new1> cat sitting in a bucket",
    num_inference_steps=100,
    guidance_scale=6.0,
    eta=1.0,
).images[0]
image.save("cat.png")

Các bước tiếp theo

Chúc mừng bạn đã huấn luyện một mô hình với Custom Diffusion! 🎉 Để tìm hiểu thêm:

Update on GitHub

←LoRA Chưng cất nhất quán tiềm ẩn→