DreamBooth

DreamBooth là một kỹ thuật huấn luyện cập nhật toàn bộ mô hình khuếch tán bằng cách huấn luyện trên chỉ một vài hình ảnh của một chủ thể hoặc phong cách. Nó hoạt động bằng cách liên kết một từ đặc biệt trong prompt với các hình ảnh mẫu.

Để tải checkpoint đã huấn luyện cho suy luận, xem Tải mô hình DreamBooth cho suy luận.

Nếu bạn đang huấn luyện trên GPU có vRAM hạn chế, bạn nên thử bật các tham số gradient_checkpointing và mixed_precision trong lệnh huấn luyện. Bạn cũng có thể giảm lượng bộ nhớ sử dụng bằng cách dùng attention tiết kiệm bộ nhớ với xFormers.

Hướng dẫn này sẽ khám phá script train_dreambooth.py để giúp bạn làm quen hơn với nó, và 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 cho script bạn đang sử dụng:

cd examples/dreambooth
pip install -r requirements.txt

🤗 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

DreamBooth rất nhạy cảm với các siêu tham số huấn luyện và dễ bị overfit. Hãy đọc bài viết trên blog Training Stable Diffusion with Dreambooth using 🧨 Diffusers để xem các cài đặt được khuyến nghị cho các chủ thể khác nhau, giúp bạn chọn được siêu tham số phù hợp.

Tệp kịch bản huấn luyện cung cấp nhiều tham số để tùy chỉnh phiên huấn luyện của bạn. Tất cả các tham số và mô tả của chúng được tìm thấy trong hàm parse_args(). Các tham số được đặt với các giá trị mặc định nên hoạt động khá tốt ngay khi sử dụng, nhưng bạn cũng có thể đặt giá trị riêng của mình trong lệnh huấn luyện nếu muốn.

Ví dụ, để huấn luyện ở định dạng bf16:

accelerate launch train_dreambooth.py \
    --mixed_precision="bf16"

Một số tham số cơ bản và quan trọng cần biết và chỉ định là:

--pretrained_model_name_or_path: tên của model trên Hub hoặc đường dẫn cục bộ đến model được huấn luyện trước
--instance_data_dir: đường dẫn đến thư mục chứa bộ dữ liệu huấn luyện (các hình ảnh ví dụ)
--instance_prompt: câu lệnh văn bản chứa từ đặc biệt cho các hình ảnh ví dụ
--train_text_encoder: có huấn luyện cả text encoder hay không
--output_dir: nơi lưu trữ mô hình đã được huấn luyện
--push_to_hub: có đẩy mô hình đã huấn luyện lên Hub hay không
--checkpointing_steps: tần suất lưu checkpoint trong quá trình huấn luyện mô hình; điều này hữu ích nếu vì một lý do nào đó quá trình huấn luyện bị gián đoạn, bạn có thể tiếp tục huấn luyện từ checkpoint đó bằng cách thêm --resume_from_checkpoint vào lệnh huấn luyện của bạn

Trọng số Min-SNR

Chiến lược trọng số Min-SNR có thể hỗ trợ quá trình huấn luyện bằng cách cân bằng lại loss để đạt được sự hội tụ nhanh hơn. Script huấn luyện hỗ trợ dự đoán epsilon (noise) hoặc v_prediction, nhưng Min-SNR tương thích với cả hai loại dự đoán. Chiến lược trọng số này chỉ được hỗ trợ bởi PyTorch.

Thêm tham số --snr_gamma và đặt nó bằng giá trị khuyến nghị là 5.0:

accelerate launch train_dreambooth.py \
  --snr_gamma=5.0

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.

--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
--class_data_dir: đường dẫn đến thư mục chứa các hình ảnh mẫu lớp đã được tạo
--class_prompt: câu lệnh văn bản mô tả lớp của các hình ảnh mẫu được tạo

accelerate launch train_dreambooth.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="path/to/class/images" \
  --class_prompt="text prompt describing class"

Huấn luyện text encoder

Để cải thiện chất lượng của các đầu ra được tạo ra, bạn cũng có thể huấn luyện text encoder bên cạnh UNet. Điều này đòi hỏi thêm bộ nhớ và bạn sẽ cần một GPU có ít nhất 24GB vRAM. Nếu bạn có phần cứng cần thiết, việc huấn luyện text encoder sẽ cho kết quả tốt hơn, đặc biệt khi tạo ảnh khuôn mặt. Bật tùy chọn này bằng cách:

accelerate launch train_dreambooth.py \
  --train_text_encoder

Script huấn luyện

DreamBooth đi kèm với các lớp dữ liệu riêng:

DreamBoothDataset: tiền xử lý các hình ảnh và hình ảnh lớp, và mã hóa các prompt để huấn luyện
PromptDataset: tạo các prompt embeddings để tạo ảnh lớp

Nếu bạn đã bật mất mát bảo tồn trước đó, các hình ảnh lớp sẽ được tạo tại đây:

sample_dataset = PromptDataset(args.class_prompt, num_new_images)
sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)

sample_dataloader = accelerator.prepare(sample_dataloader)
pipeline.to(accelerator.device)

for example in tqdm(
    sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
):
    images = pipeline(example["prompt"]).images

Tiếp theo là hàm main() xử lý việc thiết lập bộ dữ liệu cho quá trình huấn luyện và chính vòng lặp huấn luyện. Script sẽ nạp tokenizer, trình lập lịch và các mô hình:

# Load the tokenizer
if args.tokenizer_name:
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, revision=args.revision, use_fast=False)
elif args.pretrained_model_name_or_path:
    tokenizer = AutoTokenizer.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="tokenizer",
        revision=args.revision,
        use_fast=False,
    )

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = text_encoder_cls.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)

if model_has_vae(args):
    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision
    )
else:
    vae = None

unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

Sau đó, đến lúc tạo bộ dữ liệu huấn luyện và DataLoader từ DreamBoothDataset:

train_dataset = DreamBoothDataset(
    instance_data_root=args.instance_data_dir,
    instance_prompt=args.instance_prompt,
    class_data_root=args.class_data_dir if args.with_prior_preservation else None,
    class_prompt=args.class_prompt,
    class_num=args.num_class_images,
    tokenizer=tokenizer,
    size=args.resolution,
    center_crop=args.center_crop,
    encoder_hidden_states=pre_computed_encoder_hidden_states,
    class_prompt_encoder_hidden_states=pre_computed_class_prompt_encoder_hidden_states,
    tokenizer_max_length=args.tokenizer_max_length,
)

train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=args.train_batch_size,
    shuffle=True,
    collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
    num_workers=args.dataloader_num_workers,
)

Cuối cùng, vòng lặp huấn luyện sẽ xử lý các bước còn lại như chuyển đổi ảnh sang không gian ẩn, thêm nhiễu vào đầu vào, dự đoán phần dư nhiễu và tính toán tổn thất.

Nếu bạn muốn tìm hiểu thêm về cách hoạt động của vòng lặp huấn luyện, hãy xem hướng dẫn Hiểu về pipeline, mô hình và trình lập lịch, nơi phân tích chi tiết mẫu cơ bản của quá trình khử nhiễu.

Chạy script

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 một số hình ảnh của một con chó và lưu chúng vào một thư mục. Nhưng hãy nhớ rằng bạn 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).

from huggingface_hub import snapshot_download

local_dir = "./dog"
snapshot_download(
    "diffusers/dog-example",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

Đặ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 chó, và OUTPUT_DIR thành nơi bạn muốn lưu model. Bạn sẽ sử dụng sks làm từ đặc biệt để gắn kết quá trình huấn luyện.

Nếu bạn quan tâm đến việc theo dõi quá trình huấn luyện, bạn có thể định kỳ lưu các hình ảnh được tạo ra khi quá trình huấn luyện diễn ra. Thêm các tham số sau vào lệnh huấn luyện:

--validation_prompt="a photo of a sks dog"
--num_validation_images=4
--validation_steps=100

Một điều nữa trước khi bạn chạy script! Tùy thuộc vào GPU bạn đang dùng, bạn có thể cần bật một số tối ưu hóa nhất định để huấn luyện DreamBooth.

16GB

12GB

8GB

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export INSTANCE_DIR="./dog"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=400 \
  --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 vừa được huấn luyện để suy luận!

Không thể chờ để thử mô hình của bạn cho suy luận trước khi huấn luyện hoàn tất? 🤭 Hãy đảm bảo bạn đã cài đặt phiên bản mới nhất của 🤗 Accelerate.

from diffusers import DiffusionPipeline, UNet2DConditionModel
from transformers import CLIPTextModel
import torch

unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet")

# if you have trained with `--args.train_text_encoder` make sure to also load the text encoder
text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder")

pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16,
).to("cuda")

image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path_to_saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

LoRA

LoRA là một kỹ thuật huấn luyện giúp giảm đáng kể số lượng tham số có thể huấn luyện. Nhờ đó, quá trình huấn luyện diễn ra nhanh hơn và việc lưu trữ các trọng số kết quả cũng dễ dàng hơn vì kích thước của chúng nhỏ hơn rất nhiều (~100MB). Hãy sử dụng script train_dreambooth_lora.py để huấn luyện với LoRA.

Tài liệu hướng dẫn Đào tạo LoRA thảo luận chi tiết hơn về script đào tạo LoRA.

Stable Diffusion XL

Stable Diffusion XL (SDXL) là một mô hình text-to-image mạnh mẽ, tạo ra các hình ảnh độ phân giải cao, và nó bổ sung một text-encoder thứ hai vào kiến trúc của mình. Sử dụng script train_dreambooth_lora_sdxl.py để huấn luyện một mô hình SDXL với LoRA.

Tài liệu hướng dẫn huấn luyện SDXL thảo luận chi tiết hơn về script huấn luyện SDXL.

DeepFloyd IF

DeepFloyd IF là mô hình khuếch tán pixel dạng liên tiếp với ba giai đoạn. Giai đoạn đầu tiên tạo ra ảnh cơ sở, và hai giai đoạn thứ hai và thứ ba lần lượt phóng to ảnh cơ sở thành ảnh độ phân giải cao 1024x1024. Sử dụng các script train_dreambooth_lora.py hoặc train_dreambooth.py để huấn luyện mô hình DeepFloyd IF với LoRA hoặc toàn bộ mô hình.

DeepFloyd IF sử dụng phương sai dự đoán, nhưng các script huấn luyện của Diffusers sử dụng lỗi dự đoán nên các mô hình DeepFloyd IF đã được huấn luyện sẽ được chuyển sang lịch trình phương sai cố định. Các script huấn luyện sẽ tự động cập nhật cấu hình scheduler của mô hình đã được huấn luyện hoàn chỉnh cho bạn. Tuy nhiên, khi bạn tải trọng số LoRA đã lưu, bạn cũng phải cập nhật cấu hình scheduler của pipeline.

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", use_safetensors=True)

pipe.load_lora_weights("<lora weights path>")

# Update scheduler config to fixed variance schedule
pipe.scheduler = pipe.scheduler.__class__.from_config(pipe.scheduler.config, variance_type="fixed_small")

Mô hình giai đoạn 2 cần thêm các hình ảnh xác thực để tăng độ phân giải. Bạn có thể tải xuống và sử dụng một phiên bản thu nhỏ của các hình ảnh huấn luyện cho mục đích này.

from huggingface_hub import snapshot_download

local_dir = "./dog_downsized"
snapshot_download(
    "diffusers/dog-example-downsized",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

Các ví dụ mã dưới đây cung cấp cái nhìn tổng quan ngắn gọn về cách huấn luyện một mô hình DeepFloyd IF bằng cách kết hợp DreamBooth và LoRA. Một số tham số quan trọng cần lưu ý là:

--resolution=64, độ phân giải nhỏ hơn nhiều được yêu cầu vì DeepFloyd IF là mô hình khuếch tán pixel và để làm việc trên các pixel chưa nén, các hình ảnh đầu vào phải nhỏ hơn
--pre_compute_text_embeddings, tính toán trước các text embeddings để tiết kiệm bộ nhớ vì T5Model có thể chiếm rất nhiều bộ nhớ
--tokenizer_max_length=77, bạn có thể sử dụng độ dài văn bản mặc định dài hơn với T5 làm bộ mã hóa văn bản, nhưng quy trình mã hóa mô hình mặc định sử dụng độ dài văn bản ngắn hơn
--text_encoder_use_attention_mask, để truyền attention mask vào text encoder

Giai đoạn 1 LoRA DreamBooth

Giai đoạn 2 LoRA DreamBooth

Giai đoạn 1 DreamBooth

Giai đoạn 2 DreamBooth

Giai đoạn 1 của quá trình huấn luyện DeepFloyd IF với LoRA và DreamBooth yêu cầu ~28GB bộ nhớ.

export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth_dog_lora"

accelerate launch train_dreambooth_lora.py \
  --report_to wandb \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a sks dog" \
  --resolution=64 \
  --train_batch_size=4 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --scale_lr \
  --max_train_steps=1200 \
  --validation_prompt="a sks dog" \
  --validation_epochs=25 \
  --checkpointing_steps=100 \
  --pre_compute_text_embeddings \
  --tokenizer_max_length=77 \
  --text_encoder_use_attention_mask

Mẹo huấn luyện

Việc huấn luyện mô hình DeepFloyd IF có thể gây khó khăn, nhưng dưới đây là một số mẹo mà chúng tôi thấy hữu ích:

LoRA là đủ để huấn luyện mô hình giai đoạn 1 vì độ phân giải thấp của mô hình khiến việc thể hiện các chi tiết tinh tế trở nên khó khăn bất kể trường hợp nào.
Đối với các đối tượng phổ biến hoặc đơn giản, bạn không nhất thiết phải tinh chỉnh bộ tăng độ phân giải. Hãy đảm bảo rằng prompt được truyền vào bộ tăng độ phân giải đã được điều chỉnh để loại bỏ token mới khỏi prompt của instance. Ví dụ, nếu prompt ở giai đoạn 1 của bạn là “a sks dog” thì prompt ở giai đoạn 2 nên là “a dog”.
Đối với các chi tiết nhỏ hơn như khuôn mặt, việc huấn luyện hoàn toàn bộ upscaler giai đoạn 2 sẽ hiệu quả hơn so với việc huấn luyện mô hình giai đoạn 2 bằng LoRA. Đồng thời, việc sử dụng tốc độ học thấp hơn với kích thước lô lớn hơn cũng giúp ích.
Nên sử dụng tốc độ học thấp hơn để huấn luyện mô hình giai đoạn 2.
DDPMScheduler hoạt động tốt hơn so với DPMSolver được sử dụng trong các script huấn luyện.

Các bước tiếp theo

Chúc mừng bạn đã huấn luyện thành công mô hình DreamBooth! Để tìm hiểu thêm về cách sử dụng mô hình mới của bạn, hướng dẫn sau có thể hữu ích:

Tìm hiểu cách tải checkpoint DreamBooth để suy luận.
Nếu bạn đã huấn luyện với LoRA, hãy xem LoRA cho suy luận.

Update on GitHub