Skip to content

Wuerstchen

Wuerstchen 模型通过将潜在空间压缩 42 倍,大幅降低了计算成本,同时没有影响图像质量,并加速了推理。在训练过程中,Wuerstchen 使用两个模型(VQGAN + 自编码器)来压缩潜在空间,然后第三个模型(文本条件潜在扩散模型)以这种高度压缩的空间为条件来生成图像。

为了将先验模型放入 GPU 内存并加速训练,请尝试分别启用 gradient_accumulation_stepsgradient_checkpointingmixed_precision

本指南探讨了 train_text_to_image_prior.py 脚本,以帮助你更熟悉它,以及如何将其应用于你自己的用例。

在运行脚本之前,请确保你从源代码安装了库:

bash
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

然后导航到包含训练脚本的示例文件夹,并安装你正在使用的脚本所需的依赖项:

bash
cd examples/wuerstchen/text_to_image
pip install -r requirements.txt

初始化一个 🤗 Accelerate 环境:

bash
accelerate config

要设置一个默认的 🤗 Accelerate 环境,无需选择任何配置:

bash
accelerate config default

或者,如果你的环境不支持交互式 shell,比如笔记本,你可以使用:

py
from accelerate.utils import write_basic_config

write_basic_config()

脚本参数

训练脚本提供了许多参数,帮助你自定义训练运行。所有参数及其描述都可以在 parse_args() 函数中找到。它为每个参数提供了默认值,例如训练批次大小和学习率,但你也可以在训练命令中设置自己的值。

例如,要使用 fp16 格式的混合精度加速训练,请将 --mixed_precision 参数添加到训练命令中:

bash
accelerate launch train_text_to_image_prior.py \
  --mixed_precision="fp16"

大多数参数与文本到图像训练指南中的参数相同,所以让我们直接进入 Wuerstchen 训练脚本!

训练脚本

训练脚本也类似于文本到图像训练指南,但它已被修改以支持 Wuerstchen。本指南重点介绍 Wuerstchen 训练脚本中独有的代码。

main() 函数首先初始化图像编码器 - 一个EfficientNet - 以及通常的调度器和分词器。

py
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] 模型进行优化。

py
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,
)

接下来,你将对图像应用一些变换标记标题:

py
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]模型预测噪声残差。

py
pred_noise = prior(noisy_latents, timesteps, prompt_embeds)

如果你想了解更多关于训练循环的工作原理,请查看理解管道、模型和调度器教程,该教程分解了去噪过程的基本模式。

启动脚本

完成所有更改或对默认配置感到满意后,就可以启动训练脚本了!🚀

DATASET_NAME环境变量设置为来自 Hub 的数据集名称。本指南使用Naruto BLIP 标题数据集,但你也可以创建和训练自己的数据集(请参阅创建用于训练的数据集指南)。

bash
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"

训练完成后,你就可以使用新训练的模型进行推理了!

py
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 文档,了解如何使用该管道进行文本到图像的生成及其局限性。