Kandinsky 2.2

Script này 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.

Kandinsky 2.2 là một mô hình tạo ảnh từ văn bản đa ngôn ngữ, có khả năng tạo ra các hình ảnh chân thực hơn. Mô hình bao gồm một mô hình prior cho ảnh để tạo các embedding ảnh từ các prompt văn bản, và một mô hình decoder để tạo ảnh dựa trên các embedding của mô hình prior. Đó là lý do bạn sẽ tìm thấy hai script riêng biệt trong Diffusers cho Kandinsky 2.2, một để huấn luyện mô hình prior và một để huấn luyện mô hình decoder. Bạn có thể huấn luyện cả hai mô hình riêng biệt, nhưng để đạt được kết quả tốt nhất, bạn nên huấn luyện cả mô hình prior và mô hình decoder.

Tùy thuộc vào GPU của bạn, bạn có thể cần bật gradient_checkpointing (⚠️ không được hỗ trợ cho model trước!), mixed_precision và gradient_accumulation_steps để giúp model vừa với bộ nhớ và tăng tốc quá trình huấn luyện. Bạn có thể giảm mức sử dụng bộ nhớ xuống thấp hơn nữa bằng cách bật attention tiết kiệm bộ nhớ với xFormers (phiên bản v0.0.16 gặp lỗi khi huấn luyện trên một số GPU nên bạn có thể cần cài đặt phiên bản phát triển thay thế).

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

Trước khi chạy các 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/kandinsky2_2/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.

Các phần sau đây nêu bật những phần quan trọng trong các script huấn luyện để hiểu cách chỉnh sửa chúng, nhưng không đi sâu vào mọi khía cạnh của các script. Nếu bạn muốn tìm hiểu thêm, hãy tự do đọc qua các 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

Các 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(). Các script huấn luyện cung cấp các giá trị mặc định cho từng tham số, chẳng hạn như kích thước batch 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_prior.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 hãy đi thẳng vào phần hướng dẫn chi tiết các script huấn luyện Kandinsky!

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_prior.py \
  --snr_gamma=5.0

Script huấn luyện

Script huấn luyện cũng tương tự như hướng dẫn huấn luyện Text-to-image, nhưng đã được chỉnh sửa để hỗ trợ huấn luyện các mô hình prior và decoder. Hướng dẫn này tập trung vào phần mã độc quyền trong các script huấn luyện Kandinsky 2.2.

noise_scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", prediction_type="sample")
image_processor = CLIPImageProcessor.from_pretrained(
    args.pretrained_prior_model_name_or_path, subfolder="image_processor"
)
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="tokenizer")

with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
    text_encoder = CLIPTextModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="text_encoder", dtype=weight_dtype
    ).eval()

prior = PriorTransformer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
prior.train()
optimizer = optimizer_cls(
    prior.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
    return examples

model_pred = prior(
    noisy_latents,
    timestep=timesteps,
    proj_embedding=prompt_embeds,
    encoder_hidden_states=text_encoder_hidden_states,
    attention_mask=text_mask,
).predicted_image_embedding

with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    vae = VQModel.from_pretrained(
        args.pretrained_decoder_model_name_or_path, subfolder="movq", dtype=weight_dtype
    ).eval()
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
unet = UNet2DConditionModel.from_pretrained(args.pretrained_decoder_model_name_or_path, subfolder="unet")

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["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    return examples

model_pred = unet(noisy_latents, timesteps, None, added_cond_kwargs=added_cond_kwargs).sample[:, :4]

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

Bạn sẽ 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 mình, nhưng bạn cũng có thể tạo và huấn luyện trên bộ dữ liệu của riêng mình bằng cách làm theo hướng dẫn Create a dataset for training. Đặt biến môi trường DATASET_NAME thành tên của bộ dữ liệu trên Hub hoặc nếu bạn đang huấn luyện trên các tệp của riêng mình, hãy đặt biến môi trường TRAIN_DIR thành đường dẫn đến bộ dữ liệu của bạn.

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.

Để theo dõi tiến độ huấn luyện với Weights & Biases, hãy thêm tham số --report_to=wandb vào lệnh huấn luyện. Bạn cũng cần thêm --validation_prompt vào lệnh huấn luyện để theo dõi kết quả. Điều này thực sự hữu ích để gỡ lỗi mô hình và xem các kết quả trung gian.

export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_prior.py \
  --dataset_name=$DATASET_NAME \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --checkpoints_total_limit=3 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --validation_prompts="A robot naruto, 4k photo" \
  --report_to="wandb" \
  --push_to_hub \
  --output_dir="kandi2-prior-naruto-model"

export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_decoder.py \
  --dataset_name=$DATASET_NAME \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --checkpoints_total_limit=3 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --validation_prompts="A robot naruto, 4k photo" \
  --report_to="wandb" \
  --push_to_hub \
  --output_dir="kandi2-decoder-naruto-model"

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 AutoPipelineForText2Image, DiffusionPipeline
import torch

prior_pipeline = DiffusionPipeline.from_pretrained(output_dir, dtype=torch.float16)
prior_components = {"prior_" + k: v for k,v in prior_pipeline.components.items()}
pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", **prior_components, dtype=torch.float16)

pipe.enable_model_cpu_offload()
prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt, negative_prompt=negative_prompt).images[0]

from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained("path/to/saved/model", dtype=torch.float16)
pipeline.enable_model_cpu_offload()

prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt).images[0]

from diffusers import AutoPipelineForText2Image, UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained("path/to/saved/model" + "/checkpoint-<N>/unet")

pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", unet=unet, dtype=torch.float16)
pipeline.enable_model_cpu_offload()

image = pipeline(prompt="A robot naruto, 4k photo").images[0]

Các bước tiếp theo

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

• Đọc hướng dẫn Kandinsky để tìm hiểu cách sử dụng nó cho nhiều tác vụ khác nhau (text-to-image, image-to-image, inpainting, interpolation), và cách kết hợp nó với ControlNet.
• Xem hướng dẫn huấn luyện DreamBooth và LoRA để tìm hiểu cách huấn luyện một mô hình Kandinsky cá nhân hóa chỉ với một vài hình ảnh mẫu. Hai kỹ thuật huấn luyện này thậm chí có thể được kết hợp với nhau!