Văn bản sang hình ảnh

Script text-to-image mang tính thử nghiệm, và rất dễ bị overfit cũng như gặp các vấn đề như catastrophic forgetting. Hãy thử khám phá các hyperparameter khác nhau để đạt được kết quả tốt nhất trên dataset của bạn.

Các mô hình text-to-image như Stable Diffusion được điều kiện để tạo ra hình ảnh dựa trên một prompt văn bản.

Việc huấn luyện một mô hình có thể gây áp lực lớn lên phần cứng của bạn, nhưng nếu bạn bật gradient_checkpointing và mixed_precision, bạn hoàn toàn có thể huấn luyện mô hình trên một GPU 24GB duy nhất. Nếu bạn đang huấn luyện với kích thước batch lớn hơn hoặc muốn huấn luyện nhanh hơn, tốt hơn hết là nên sử dụng các GPU có bộ nhớ lớn hơn 30GB. Bạn có thể giảm lượng bộ nhớ sử dụng bằng cách bật cơ chế attention tiết kiệm bộ nhớ với xFormers.

Hướng dẫn này sẽ khám phá script huấn luyện train_text_to_image.py để giúp bạn làm quen 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/text_to_image
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.

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_text_to_image.py \
  --mixed_precision="fp16"

Một số tham số cơ bản và quan trọng bao gồm:

• --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
• --dataset_name: tên của bộ dữ liệu trên Hub hoặc đường dẫn cục bộ đến bộ dữ liệu để huấn luyện
• --image_column: tên của cột hình ảnh trong bộ dữ liệu để huấn luyện
• --caption_column: tên của cột văn bản trong bộ dữ liệu để huấn luyện
• --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_text_to_image.py \
  --snr_gamma=5.0

Bạn có thể so sánh các bề mặt loss cho các giá trị snr_gamma khác nhau trong báo cáo Weights and Biases này. Đối với các bộ dữ liệu nhỏ hơn, tác động của Min-SNR có thể không rõ ràng bằng so với các bộ dữ liệu lớn hơn.

Script huấn luyện

Mã tiền xử lý tập dữ liệu và vòng lặp huấn luyện nằm trong hàm main(). Nếu bạn cần điều chỉnh script huấn luyện, đây là nơi bạn cần thực hiện các thay đổi.

Script train_text_to_image bắt đầu bằng việc tải scheduler và tokenizer. Bạn có thể chọn sử dụng một scheduler khác tại đây nếu muốn:

noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
)

Sau đó, script tải mô hình UNet:

load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
model.register_to_config(**load_model.config)

model.load_state_dict(load_model.state_dict())

Tiếp theo, các cột văn bản và hình ảnh của tập dữ liệu cần được tiền xử lý. Hàm tokenize_captions xử lý việc token hóa các đầu vào, và hàm train_transforms chỉ định loại biến đổi cần áp dụng cho hình ảnh. Cả hai hàm này được đóng gói trong preprocess_train:

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["pixel_values"] = [train_transforms(image) for image in images]
    examples["input_ids"] = tokenize_captions(examples)
    return examples

Cuối cùng, vòng lặp huấn luyện xử lý mọi thứ còn lại. Nó mã hóa ảnh vào không gian tiềm ẩn, thêm nhiễu vào các giá trị tiềm ẩn, tính toán các nhúng văn bản để điều kiện hóa, cập nhật tham số mô hình, và lưu cũng như đẩy mô hình lên Hub. Nếu bạn muốn tìm hiểu thêm về cách vòng lặp huấn luyện hoạt động, hãy xem hướng dẫn Hiểu về pipeline, mô hình và trình lập lịch phân tích chi tiết mẫu cơ bản của quá trình khử nhiễu.

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! 🚀

Hãy huấn luyện trên bộ dữ liệu Naruto BLIP captions để tạo ra các nhân vật Naruto của riêng bạn. Đặt các biến môi trường MODEL_NAME và dataset_name thành tên mô hình và tên bộ dữ liệu (từ Hub hoặc đường dẫn cục bộ). Nếu bạn đang huấn luyện trên nhiều hơn một GPU, hãy thêm tham số --multi_gpu vào lệnh accelerate launch.

Để huấn luyện trên một bộ dữ liệu cục bộ, hãy đặt các biến môi trường TRAIN_DIR và OUTPUT_DIR thành đường dẫn của bộ dữ liệu và nơi lưu trữ mô hình.

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export dataset_name="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --enable_xformers_memory_efficient_attention \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model" \
  --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:

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("path/to/saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

image = pipeline(prompt="yoda").images[0]
image.save("yoda-naruto.png")

Các bước tiếp theo

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

• Tìm hiểu cách tải trọng số LoRA để suy luận nếu bạn đã huấn luyện mô hình của mình với LoRA.
• Tìm hiểu thêm về cách các tham số nhất định như guidance scale hoặc các kỹ thuật như prompt weighting có thể giúp bạn kiểm soát suy luận trong hướng dẫn tác vụ Text-to-image.