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

Mô hình nhất quán tiềm ẩn (LCM) có khả năng tạo ra hình ảnh chất lượng cao chỉ trong vài bước, đánh dấu một bước tiến lớn vì nhiều pipeline yêu cầu ít nhất 25+ bước. LCM được tạo ra bằng cách áp dụng phương pháp chưng cất nhất quán tiềm ẩn lên bất kỳ mô hình Stable Diffusion nào. Phương pháp này hoạt động bằng cách áp dụng chưng cất có hướng dẫn một giai đoạn lên không gian tiềm ẩn, và tích hợp phương pháp bỏ qua bước để nhất quán bỏ qua các bước thời gian nhằm tăng tốc quá trình chưng cất (xem mục 4.1, 4.2 và 4.3 của bài báo để biết thêm chi tiết).

Nếu bạn đang huấn luyện trên GPU có vRAM hạn chế, hãy thử bật gradient_checkpointing, gradient_accumulation_steps và mixed_precision để giảm mức sử dụng bộ nhớ và tăng tốc độ huấn luyện. Bạn có thể giảm mức sử dụng bộ nhớ còn nhiều hơn nữa bằng cách bật cơ chế attention tiết kiệm bộ nhớ với xFormers và trình tối ưu hóa 8-bit của bitsandbytes.

Hướng dẫn này sẽ khám phá script train_lcm_distill_sd_wds.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 .

Sau đó, điều hướng đến thư mục ví dụ chứa 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/consistency_distillation
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 (thử bật torch.compile để tăng tốc đáng kể quá trình huấn luyện):

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.

Tham số của script

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.

Tệp script huấn luyện cung cấp nhiều tham số để giúp bạn tùy chỉnh phiên huấn luyện của mình. Tất cả các tham số và mô tả của chúng được tìm thấy trong hàm parse_args(). Hàm này cung cấp các giá trị mặc định cho từng tham số, chẳng hạn như kích thước lô huấn luyện và tốc độ học, 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ụ, để tăng tốc độ huấn luyện với độ chính xác hỗn hợp sử dụng định dạng fp16, hãy thêm tham số --mixed_precision vào lệnh huấn luyện:

accelerate launch train_lcm_distill_sd_wds.py \
  --mixed_precision="fp16"

Hầu hết các tham số đều giống với các tham số trong hướng dẫn huấn luyện Text-to-image, vì vậy bạn sẽ tập trung vào các tham số liên quan đến latent consistency distillation trong hướng dẫn này.

--pretrained_teacher_model: đường dẫn đến mô hình khuếch tán tiềm ẩn (latent diffusion) đã được huấn luyện trước để sử dụng làm mô hình giáo viên
--pretrained_vae_model_name_or_path: đường dẫn đến một VAE đã được huấn luyện trước; VAE của SDXL được biết là gặp phải sự bất ổn định về mặt số học, vì vậy tham số này cho phép bạn chỉ định một VAE thay thế (như VAE này) do madebyollin tạo ra, hoạt động ở fp16)
--w_min và --w_max: giá trị tối thiểu và tối đa của guidance scale cho việc lấy mẫu guidance scale
--num_ddim_timesteps: số lượng bước thời gian cho việc lấy mẫu DDIM
--loss_type: loại loss (L2 hoặc Huber) được tính cho latent consistency distillation; loss Huber thường được ưu tiên vì nó bền vững hơn với các giá trị ngoại lai
--huber_c: tham số của hàm mất Huber

Script huấn luyện

Tệp kịch bản huấn luyện bắt đầu bằng việc tạo một lớp dữ liệu - Text2ImageDataset - để tiền xử lý các hình ảnh và tạo tập dữ liệu huấn luyện.

def transform(example):
    image = example["image"]
    image = TF.resize(image, resolution, interpolation=transforms.InterpolationMode.BILINEAR)

    c_top, c_left, _, _ = transforms.RandomCrop.get_params(image, output_size=(resolution, resolution))
    image = TF.crop(image, c_top, c_left, resolution, resolution)
    image = TF.to_tensor(image)
    image = TF.normalize(image, [0.5], [0.5])

    example["image"] = image
    return example

Để cải thiện hiệu suất khi đọc và ghi các bộ dữ liệu lớn được lưu trữ trên đám mây, script này sử dụng định dạng WebDataset để tạo một pipeline tiền xử lý nhằm áp dụng các phép biến đổi và tạo ra một dataset cùng dataloader cho quá trình huấn luyện. Các hình ảnh được xử lý và đưa vào vòng lặp huấn luyện mà không cần phải tải toàn bộ bộ dữ liệu trước.

processing_pipeline = [
    wds.decode("pil", handler=wds.ignore_and_continue),
    wds.rename(image="jpg;png;jpeg;webp", text="text;txt;caption", handler=wds.warn_and_continue),
    wds.map(filter_keys({"image", "text"})),
    wds.map(transform),
    wds.to_tuple("image", "text"),
]

Trong hàm main(), tất cả các thành phần cần thiết như noise scheduler, tokenizers, text encoders và VAE đều được nạp. Teacher UNet cũng được nạp tại đây và sau đó bạn có thể tạo student UNet từ teacher UNet. Student UNet được cập nhật bởi optimizer trong quá trình huấn luyện.

teacher_unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_teacher_model, subfolder="unet", revision=args.teacher_revision
)

unet = UNet2DConditionModel(**teacher_unet.config)
unet.load_state_dict(teacher_unet.state_dict(), strict=False)
unet.train()

Bây giờ bạn có thể tạo optimizer để cập nhật các tham số của UNet:

optimizer = optimizer_class(
    unet.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Tạo dataset:

dataset = Text2ImageDataset(
    train_shards_path_or_url=args.train_shards_path_or_url,
    num_train_examples=args.max_train_samples,
    per_gpu_batch_size=args.train_batch_size,
    global_batch_size=args.train_batch_size * accelerator.num_processes,
    num_workers=args.dataloader_num_workers,
    resolution=args.resolution,
    shuffle_buffer_size=1000,
    pin_memory=True,
    persistent_workers=True,
)
train_dataloader = dataset.train_dataloader

Tiếp theo, bạn đã sẵn sàng để thiết lập vòng lặp huấn luyện và triển khai phương pháp chưng cất tính nhất quán tiềm ẩn (xem Thuật toán 1 trong bài báo để biết thêm chi tiết). Phần này của script chịu trách nhiệm thêm nhiễu vào các biến tiềm ẩn, lấy mẫu và tạo nhúng thang độ dẫn hướng, cũng như dự đoán ảnh gốc từ nhiễu.

pred_x_0 = predicted_origin(
    noise_pred,
    start_timesteps,
    noisy_model_input,
    noise_scheduler.config.prediction_type,
    alpha_schedule,
    sigma_schedule,
)

model_pred = c_skip_start * noisy_model_input + c_out_start * pred_x_0

Tiếp theo, nó lấy dự đoán của mô hình teacher và dự đoán của LCM, tính toán loss, sau đó backpropagate loss đó vào LCM.

if args.loss_type == "l2":
    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
elif args.loss_type == "huber":
    loss = torch.mean(
        torch.sqrt((model_pred.float() - target.float()) ** 2 + args.huber_c**2) - args.huber_c
    )

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 tutorial về cách hiểu các 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ây giờ bạn đã sẵn sàng để chạy script huấn luyện và bắt đầu quá trình chưng cất!

Trong hướng dẫn này, bạn sẽ sử dụng --train_shards_path_or_url để chỉ định đường dẫn đến bộ dữ liệu Conceptual Captions 12M được lưu trên Hub tại đây. Đặt biến môi trường MODEL_DIR thành tên của mô hình teacher và OUTPUT_DIR thành nơi bạn muốn lưu mô hình.

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/saved/model"

accelerate launch train_lcm_distill_sd_wds.py \
    --pretrained_teacher_model=$MODEL_DIR \
    --output_dir=$OUTPUT_DIR \
    --mixed_precision=fp16 \
    --resolution=512 \
    --learning_rate=1e-6 --loss_type="huber" --ema_decay=0.95 --adam_weight_decay=0.0 \
    --max_train_steps=1000 \
    --max_train_samples=4000000 \
    --dataloader_num_workers=8 \
    --train_shards_path_or_url="pipe:curl -L -s https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset/resolve/main/data/{00000..01099}.tar?download=true" \
    --validation_steps=200 \
    --checkpointing_steps=200 --checkpoints_total_limit=10 \
    --train_batch_size=12 \
    --gradient_checkpointing --enable_xformers_memory_efficient_attention \
    --gradient_accumulation_steps=1 \
    --use_8bit_adam \
    --resume_from_checkpoint=latest \
    --report_to=wandb \
    --seed=453645634 \
    --push_to_hub

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

from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
import torch

unet = UNet2DConditionModel.from_pretrained("your-username/your-model", dtype=torch.float16, variant="fp16")
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, dtype=torch.float16, variant="fp16")

pipeline.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "sushi rolls in the form of panda heads, sushi platter"

image = pipeline(prompt, num_inference_steps=4, guidance_scale=1.0).images[0]

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 nhiều (~100MB). Hãy sử dụng script train_lcm_distill_lora_sd_wds.py hoặc train_lcm_distill_lora_sdxl.wds.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_lcm_distill_sdxl_wds.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.

Các bước tiếp theo

Chúc mừng bạn đã chưng cất một mô hình LCM! Để tìm hiểu thêm về LCM, các tài liệu sau có thể hữu ích:

Tìm hiểu cách sử dụng LCMs để suy luận cho text-to-image, image-to-image và với các checkpoint LoRA.
Đọc bài đăng trên blog SDXL trong 4 bước với Latent Consistency LoRAs để tìm hiểu thêm về SDXL LCM-LoRA cho suy luận siêu nhanh, so sánh chất lượng, chuẩn đánh hiệu năng và nhiều hơn nữa.

Update on GitHub

Diffusers

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

Tham số của script

Script huấn luyện

Chạy script

LoRA

Stable Diffusion XL

Các bước tiếp theo