Diffusers 文件
DreamBooth
並獲得增強的文件體驗
開始使用
DreamBooth
DreamBooth 是一種訓練技術，透過僅使用少量影像（可以是特定主題或風格的影像）進行訓練來更新整個擴散模型。它透過將提示中的一個特殊詞與示例影像相關聯來實現這一目標。
如果你的GPU的vRAM有限，你應該嘗試在訓練命令中啟用`gradient_checkpointing`和`mixed_precision`引數。你還可以透過使用xFormers的高效記憶體注意力來減少記憶體佔用。JAX/Flax訓練也支援在TPU和GPU上進行高效訓練，但它不支援梯度檢查點或xFormers。如果你想用Flax更快地訓練，你的GPU應該有>30GB的記憶體。
本指南將探討 train_dreambooth.py 指令碼，幫助您熟悉它，以及如何根據您自己的用例進行調整。
在執行指令碼之前，請確保從原始碼安裝庫
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .導航到包含訓練指令碼的示例資料夾並安裝您正在使用的指令碼所需的依賴項
cd examples/dreambooth
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()最後，如果您想在自己的資料集上訓練模型，請檢視 建立訓練資料集 指南，瞭解如何建立與訓練指令碼相容的資料集。
以下部分重點介紹了訓練指令碼中對於理解如何修改它很重要的部分，但它沒有詳細介紹指令碼的每個方面。如果您有興趣瞭解更多資訊，請隨意閱讀 指令碼，如果您有任何問題或疑慮，請告訴我們。
指令碼引數
DreamBooth 對訓練超引數非常敏感，很容易過擬合。閱讀使用🧨 Diffusers 訓練 Stable Diffusion 和 Dreambooth 部落格文章，瞭解不同主題的推薦設定，以幫助您選擇合適的超引數。
該訓練指令碼提供了許多引數來自定義您的訓練執行。所有引數及其描述都可以在parse_args()函式中找到。引數設定了預設值，這些預設值應該開箱即用，但您也可以在訓練命令中設定自己的值，如果您願意的話。
例如，以 bf16 格式訓練
accelerate launch train_dreambooth.py \
--mixed_precision="bf16"一些需要了解和指定的基本且重要的引數是：
--pretrained_model_name_or_path：Hub 上的模型名稱或預訓練模型的本地路徑--instance_data_dir：包含訓練資料集（示例影像）的資料夾路徑--instance_prompt：包含示例影像特殊詞的文字提示--train_text_encoder：是否也訓練文字編碼器--output_dir：訓練好的模型的儲存位置--push_to_hub：是否將訓練好的模型推送到 Hub--checkpointing_steps：模型訓練過程中儲存檢查點的頻率；如果訓練因某種原因中斷，這很有用，您可以透過在訓練命令中新增--resume_from_checkpoint從該檢查點繼續訓練
最小信噪比加權
最小信噪比 (Min-SNR) 加權策略透過重新平衡損失以實現更快的收斂來幫助訓練。訓練指令碼支援預測 `epsilon` (噪聲) 或 `v_prediction`，但 Min-SNR 與兩種預測型別都相容。此加權策略僅受 PyTorch 支援，在 Flax 訓練指令碼中不可用。
新增 --snr_gamma 引數並將其設定為推薦值 5.0
accelerate launch train_dreambooth.py \ --snr_gamma=5.0
先驗保留損失
先驗保留損失是一種利用模型自身生成的樣本來幫助它學習如何生成更多樣化影像的方法。由於這些生成的樣本影像與您提供的影像屬於同一類別，它們有助於模型保留其已學到的類別資訊，並利用其已知的類別知識來建立新的構圖。
--with_prior_preservation: 是否使用先驗保留損失--prior_loss_weight: 控制先驗保留損失對模型的影響--class_data_dir: 包含生成類樣本影像的資料夾路徑--class_prompt: 描述生成樣本影像類別的文字提示
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"訓練文字編碼器
為了提高生成輸出的質量，除了 UNet，您還可以訓練文字編碼器。這需要額外的記憶體，並且您需要一個至少具有 24GB vRAM 的 GPU。如果您擁有必要的硬體，那麼訓練文字編碼器會產生更好的結果，尤其是在生成人臉影像時。透過以下方式啟用此選項：
accelerate launch train_dreambooth.py \ --train_text_encoder
訓練指令碼
DreamBooth 帶有自己的資料集類：
DreamBoothDataset：預處理影像和類影像，並對提示進行標記化以進行訓練PromptDataset：生成提示嵌入以生成類影像
如果您啟用了先驗保留損失，則在此處生成類影像：
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接下來是main()函式，它負責設定訓練資料集和訓練迴圈本身。指令碼載入分詞器、排程器和模型。
# 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
)然後，是時候從 DreamBoothDataset 建立訓練資料集和 DataLoader 了。
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,
)最後，訓練迴圈負責剩餘的步驟，例如將影像轉換為潛在空間，向輸入新增噪聲，預測噪聲殘差，以及計算損失。
如果您想了解更多關於訓練迴圈如何工作的資訊，請檢視 理解管道、模型和排程器 教程，它分解了去噪過程的基本模式。
啟動指令碼
您現在可以啟動訓練指令碼了！🚀
對於本指南，您將下載一些狗的影像並將其儲存在一個目錄中。但請記住，如果您願意，可以建立並使用自己的資料集（請參閱建立訓練資料集指南）。
from huggingface_hub import snapshot_download
local_dir = "./dog"
snapshot_download(
"diffusers/dog-example",
local_dir=local_dir,
repo_type="dataset",
ignore_patterns=".gitattributes",
)將環境變數 `MODEL_NAME` 設定為 Hub 上的模型 ID 或本地模型的路徑，`INSTANCE_DIR` 設定為您剛下載狗影像的路徑，`OUTPUT_DIR` 設定為您希望儲存模型的路徑。您將使用 `sks` 作為特殊詞來關聯訓練。
如果您有興趣跟蹤訓練過程，可以定期儲存生成的影像，隨著訓練的進行。將以下引數新增到訓練命令中：
--validation_prompt="a photo of a sks dog"
--num_validation_images=4
--validation_steps=100在啟動指令碼之前還有一件事！根據您擁有的 GPU，您可能需要啟用某些最佳化才能訓練 DreamBooth。
在 16GB GPU 上，您可以使用 bitsandbytes 8 位最佳化器和梯度檢查點來幫助您訓練 DreamBooth 模型。安裝 bitsandbytes
pip install bitsandbytes
然後，將以下引數新增到您的訓練命令中：
accelerate launch train_dreambooth.py \ --gradient_checkpointing \ --use_8bit_adam \
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訓練完成後，您可以使用新訓練的模型進行推理！
等不及在訓練完成前嘗試您的模型進行推理？🤭 確保您已安裝最新版本的 🤗 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", torch_dtype=torch.float16, use_safetensors=True).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")LoRA
LoRA 是一種訓練技術，可顯著減少可訓練引數的數量。因此，訓練速度更快，並且更容易儲存生成的權重，因為它們要小得多（約 100MB）。使用 train_dreambooth_lora.py 指令碼使用 LoRA 進行訓練。
LoRA 訓練指令碼在LoRA 訓練指南中進行了更詳細的討論。
Stable Diffusion XL
Stable Diffusion XL (SDXL) 是一個功能強大的文字到影像模型，可生成高解析度影像，並在其架構中添加了第二個文字編碼器。使用 train_dreambooth_lora_sdxl.py 指令碼使用 LoRA 訓練 SDXL 模型。
SDXL 訓練指令碼在 SDXL 訓練指南中有更詳細的討論。
DeepFloyd IF
DeepFloyd IF 是一個三階段級聯畫素擴散模型。第一階段生成一個基礎影像，第二和第三階段逐步將基礎影像上取樣為 1024x1024 的高解析度影像。使用 train_dreambooth_lora.py 或 train_dreambooth.py 指令碼來訓練帶 LoRA 的 DeepFloyd IF 模型或完整模型。
DeepFloyd IF 使用預測方差，但 Diffusers 訓練指令碼使用預測誤差，因此訓練後的 DeepFloyd IF 模型會切換到固定方差排程。訓練指令碼將為您更新完全訓練模型的排程器配置。但是，當您載入儲存的 LoRA 權重時，您還必須更新管道的排程器配置。
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")第二階段模型需要額外的驗證影像進行放大。您可以下載並使用訓練影像的縮小版本。
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",
)以下程式碼示例簡要概述瞭如何結合 DreamBooth 和 LoRA 訓練 DeepFloyd IF 模型。需要注意的一些重要引數是：
--resolution=64，需要更小的解析度，因為 DeepFloyd IF 是畫素擴散模型，為了處理未壓縮畫素，輸入影像必須更小--pre_compute_text_embeddings，提前計算文字嵌入以節省記憶體，因為 T5Model 會佔用大量記憶體--tokenizer_max_length=77，您可以將 T5 作為文字編碼器使用更長的預設文字長度，但預設模型編碼過程使用較短的文字長度--text_encoder_use_attention_mask，將注意力掩碼傳遞給文字編碼器
使用 LoRA 和 DreamBooth 訓練 DeepFloyd IF 的第一階段需要大約 28GB 的記憶體。
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訓練技巧
訓練 DeepFloyd IF 模型可能具有挑戰性，但以下是一些我們發現有用的技巧：
- LoRA 足以訓練第一階段模型，因為模型的低解析度使得細節表示無論如何都很困難。
- 對於常見或簡單的物件，您不一定需要微調升取樣器。請確保傳遞給升取樣器的提示已調整，以從例項提示中刪除新的標記。例如，如果您的第一階段提示是“a sks dog”，那麼您的第二階段提示應該是“a dog”。
- 對於面部等更精細的細節，完全訓練第二階段升取樣器比使用 LoRA 訓練第二階段模型更好。使用較低的學習率和較大的批次大小也有幫助。
- 應使用較低的學習率來訓練第二階段模型。
- DDPMScheduler 比訓練指令碼中使用的 DPMSolver 效果更好。
後續步驟
恭喜您訓練好您的 DreamBooth 模型！要了解如何使用您的新模型，以下指南可能會有所幫助：
- 如果您使用 LoRA 訓練了您的模型，請了解如何載入 DreamBooth 模型進行推理。