Diffusers 文档
Wuerstchen
并获取增强的文档体验
开始使用
Wuerstchen
Wuerstchen 模型通过将潜在空间压缩 42 倍，大幅降低了计算成本，同时不影响图像质量并加速了推理。在训练期间，Wuerstchen 使用两个模型（VQGAN + 自动编码器）来压缩潜在空间，然后第三个模型（文本条件潜在扩散模型）以这种高度压缩的空间为条件生成图像。
为了将先验模型放入 GPU 内存并加速训练，请尝试分别启用 gradient_accumulation_steps、gradient_checkpointing 和 mixed_precision。
本指南探讨了 train_text_to_image_prior.py 脚本，以帮助您更熟悉它，以及如何将其用于您自己的用例。
在运行脚本之前，请确保从源代码安装库
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .然后导航到包含训练脚本的示例文件夹，并安装您正在使用的脚本所需的依赖项
cd examples/wuerstchen/text_to_image
pip install -r requirements.txt🤗 Accelerate 是一个库，旨在帮助您在多个 GPU/TPU 上或使用混合精度进行训练。它将根据您的硬件和环境自动配置您的训练设置。请查看 🤗 Accelerate 快速入门 以了解更多信息。
初始化 🤗 Accelerate 环境
accelerate config
要设置默认的 🤗 Accelerate 环境而不选择任何配置
accelerate config default
或者，如果您的环境不支持交互式 shell（例如 notebook），您可以使用
from accelerate.utils import write_basic_config
write_basic_config()最后，如果您想在自己的数据集上训练模型，请查看 创建用于训练的数据集 指南，了解如何创建与训练脚本配合使用的数据集。
以下部分重点介绍了训练脚本中对于理解如何修改脚本很重要的部分，但并未详细介绍 script 的每个方面。如果您有兴趣了解更多信息，请随时阅读脚本，如果您有任何问题或疑虑，请告诉我们。
脚本参数
训练脚本提供了许多参数来帮助您自定义训练运行。所有参数及其描述都可以在 parse_args() 函数中找到。它为每个参数提供了默认值，例如训练批大小和学习率，但如果您愿意，也可以在训练命令中设置自己的值。
例如，要使用 fp16 格式的混合精度加速训练，请将 --mixed_precision 参数添加到训练命令中
accelerate launch train_text_to_image_prior.py \
--mixed_precision="fp16"大多数参数与 文本到图像 训练指南中的参数相同，因此让我们直接深入了解 Wuerstchen 训练脚本！
训练脚本
训练脚本也类似于 文本到图像 训练指南，但已进行修改以支持 Wuerstchen。本指南重点介绍 Wuerstchen 训练脚本特有的代码。
main() 函数首先初始化图像编码器 - 一个 EfficientNet - 以及常用的 scheduler 和 tokenizer。
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
pretrained_checkpoint_file = hf_hub_download("dome272/wuerstchen", filename="model_v2_stage_b.pt")
state_dict = torch.load(pretrained_checkpoint_file, map_location="cpu")
image_encoder = EfficientNetEncoder()
image_encoder.load_state_dict(state_dict["effnet_state_dict"])
image_encoder.eval()您还将加载 WuerstchenPrior 模型以进行优化。
prior = WuerstchenPrior.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
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["effnet_pixel_values"] = [effnet_transforms(image) for image in images]
examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
return examples最后，训练循环 处理使用 EfficientNetEncoder 将图像压缩到潜在空间、向潜在空间添加噪声以及使用 WuerstchenPrior 模型预测噪声残差。
pred_noise = prior(noisy_latents, timesteps, prompt_embeds)
如果您想了解有关训练循环如何工作的更多信息，请查看 理解 pipelines、模型和 schedulers 教程，其中分解了去噪过程的基本模式。
启动脚本
一旦您完成了所有更改，或者您对默认配置感到满意，就可以启动训练脚本了！🚀
将 DATASET_NAME 环境变量设置为 Hub 中的数据集名称。本指南使用 Naruto BLIP captions 数据集，但您也可以创建自己的数据集并在其上进行训练（请参阅 创建用于训练的数据集 指南）。
要使用 Weights & Biases 监控训练进度，请将 --report_to=wandb 参数添加到训练命令。您还需要将 --validation_prompt 添加到训练命令，以跟踪结果。这对于调试模型和查看中间结果非常有用。
export DATASET_NAME="lambdalabs/naruto-blip-captions"
accelerate launch train_text_to_image_prior.py \
--mixed_precision="fp16" \
--dataset_name=$DATASET_NAME \
--resolution=768 \
--train_batch_size=4 \
--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--dataloader_num_workers=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="wuerstchen-prior-naruto-model"训练完成后，您可以使用新训练的模型进行推理！
import torch
from diffusers import AutoPipelineForText2Image
from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
pipeline = AutoPipelineForText2Image.from_pretrained("path/to/saved/model", torch_dtype=torch.float16).to("cuda")
caption = "A cute bird naruto holding a shield"
images = pipeline(
caption,
width=1024,
height=1536,
prior_timesteps=DEFAULT_STAGE_C_TIMESTEPS,
prior_guidance_scale=4.0,
num_images_per_prompt=2,
).images下一步
恭喜您训练了一个 Wuerstchen 模型！要了解有关如何使用新模型的更多信息，以下内容可能会有所帮助
- 查看 Wuerstchen API 文档，以了解有关如何使用 pipeline 进行文本到图像生成及其局限性的更多信息。