Diffusers 文件
Kandinsky 2.2
並獲得增強的文件體驗
開始使用
Kandinsky 2.2
此指令碼是實驗性的,很容易過擬合併遇到災難性遺忘等問題。嘗試探索不同的超引數,以在您的資料集上獲得最佳結果。
Kandinsky 2.2 是一個多語言文字到影像模型,能夠生成更逼真的影像。該模型包括一個影像先驗模型用於從文字提示建立影像嵌入,以及一個解碼器模型用於根據先驗模型的嵌入生成影像。這就是為什麼您會在 Diffusers 中找到兩個獨立的 Kandinsky 2.2 指令碼,一個用於訓練先驗模型,一個用於訓練解碼器模型。您可以單獨訓練這兩個模型,但要獲得最佳結果,您應該同時訓練先驗模型和解碼器模型。
根據您的 GPU,您可能需要啟用 gradient_checkpointing
(⚠️ 不支援先驗模型!)、mixed_precision
和 gradient_accumulation_steps
以幫助模型適應記憶體並加速訓練。您可以透過使用 xFormers 啟用記憶體高效注意力來進一步減少記憶體使用(版本 v0.0.16 在某些 GPU 上訓練失敗,因此您可能需要安裝開發版本)。
本指南探討了 train_text_to_image_prior.py 和 train_text_to_image_decoder.py 指令碼,以幫助您更熟悉它們,以及如何將其用於您自己的用例。
在執行指令碼之前,請確保從原始碼安裝庫
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
然後導航到包含訓練指令碼的示例資料夾並安裝您正在使用的指令碼所需的依賴項
cd examples/kandinsky2_2/text_to_image
pip install -r requirements.txt
🤗 Accelerate 是一個幫助您在多個 GPU/TPU 或混合精度下進行訓練的庫。它將根據您的硬體和環境自動配置您的訓練設定。請檢視 🤗 Accelerate 快速入門 以瞭解更多資訊。
初始化 🤗 Accelerate 環境
accelerate config
要設定預設的 🤗 Accelerate 環境而不選擇任何配置
accelerate config default
或者如果您的環境不支援互動式 shell(例如筆記本),您可以使用
from accelerate.utils import write_basic_config
write_basic_config()
最後,如果您想在自己的資料集上訓練模型,請檢視 建立訓練資料集 指南,瞭解如何建立與訓練指令碼相容的資料集。
以下部分重點介紹了訓練指令碼中對於理解如何修改它很重要的部分,但它沒有詳細介紹指令碼的各個方面。如果您有興趣瞭解更多資訊,請隨時閱讀指令碼,如有任何疑問或擔憂,請告訴我們。
指令碼引數
訓練指令碼提供了許多引數來幫助您自定義訓練執行。所有引數及其描述都可以在 parse_args()
函式中找到。訓練指令碼為每個引數提供了預設值,例如訓練批次大小和學習率,但您也可以在訓練命令中設定自己的值。
例如,要使用 fp16 格式的混合精度加速訓練,請將 --mixed_precision
引數新增到訓練命令中
accelerate launch train_text_to_image_prior.py \
--mixed_precision="fp16"
大多數引數與 文字到影像 訓練指南中的引數相同,所以讓我們直接進入 Kandinsky 訓練指令碼的演練!
最小信噪比加權
最小信噪比 (Min-SNR) 加權策略可以透過重新平衡損失來實現更快收斂,從而有助於訓練。訓練指令碼支援預測 epsilon
(噪聲) 或 v_prediction
,但 Min-SNR 相容這兩種預測型別。此加權策略僅受 PyTorch 支援,在 Flax 訓練指令碼中不可用。
新增 --snr_gamma
引數並將其設定為推薦值 5.0
accelerate launch train_text_to_image_prior.py \ --snr_gamma=5.0
訓練指令碼
訓練指令碼也類似於 文字到影像 訓練指南,但它已修改為支援訓練先驗模型和解碼器模型。本指南側重於 Kandinsky 2.2 訓練指令碼獨有的程式碼。
main()
函式包含準備資料集和訓練模型的程式碼。
您會立即注意到,訓練指令碼的其中一個主要區別是,除了排程器和分詞器之外,它還載入了 CLIPImageProcessor 用於預處理影像,以及 CLIPVisionModelWithProjection 模型用於編碼影像
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", torch_dtype=weight_dtype
).eval()
text_encoder = CLIPTextModelWithProjection.from_pretrained(
args.pretrained_prior_model_name_or_path, subfolder="text_encoder", torch_dtype=weight_dtype
).eval()
Kandinsky 使用 PriorTransformer 生成影像嵌入,因此您需要設定最佳化器以學習先驗模型的引數。
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,
)
接下來,輸入字幕被分詞,影像由 CLIPImageProcessor 進行 預處理
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
如果您想了解更多關於訓練迴圈如何工作的資訊,請檢視 理解管道、模型和排程器 教程,它分解了去噪過程的基本模式。
啟動指令碼
完成所有更改或對預設配置滿意後,您就可以啟動訓練指令碼了!🚀
您將使用 Naruto BLIP captions 資料集來生成您自己的火影忍者角色,但您也可以透過遵循 建立訓練資料集 指南來建立並訓練您自己的資料集。將環境變數 DATASET_NAME
設定為 Hub 上資料集的名稱,如果您正在訓練自己的檔案,則將環境變數 TRAIN_DIR
設定為資料集的路徑。
如果您在多個 GPU 上進行訓練,請將 --multi_gpu
引數新增到 accelerate launch
命令中。
要使用 Weights & Biases 監控訓練進度,請將 --report_to=wandb
引數新增到訓練命令中。您還需要將 --validation_prompt
新增到訓練命令中以跟蹤結果。這對於除錯模型和檢視中間結果非常有用。
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"
訓練完成後,您就可以使用新訓練的模型進行推理了!
from diffusers import AutoPipelineForText2Image, DiffusionPipeline
import torch
prior_pipeline = DiffusionPipeline.from_pretrained(output_dir, torch_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, torch_dtype=torch.float16)
pipe.enable_model_cpu_offload()
prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt, negative_prompt=negative_prompt).images[0]
隨意將 kandinsky-community/kandinsky-2-2-decoder
替換為您自己訓練的解碼器檢查點!
下一步
恭喜您訓練了一個 Kandinsky 2.2 模型!要了解如何使用您的新模型,以下指南可能會有所幫助
- 閱讀 Kandinsky 指南,瞭解如何將其用於各種不同任務(文字到影像、影像到影像、影像修復、插值),以及如何將其與 ControlNet 結合使用。
- 檢視 DreamBooth 和 LoRA 訓練指南,瞭解如何僅使用少量示例影像來訓練個性化的 Kandinsky 模型。這兩種訓練技術甚至可以結合使用!