重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和一个新库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 中的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

Stable Diffusion#

本节简要概述 NeMo 框架中的 stable diffusion 模型。

模型介绍#

Stable Diffusion 是一款先进的文本到图像扩散模型，它使用大量的图像和文本对数据集进行训练。其核心功能是通过消除噪声来优化和增强图像，从而产生清晰的输出视觉效果。当模型接收到名为 z0 的图像时，会系统地注入噪声。在以 “t” 标记的每次迭代中，图像（现在称为 zt）变得越来越失真。随着 “t” 值的增加，图像越来越接近完全噪声。此外，当提供文本提示或时间步 “t” 等具体细节时，模型可以准确地确定引入到 zt 的噪声程度。

Stable diffusion 有三个主要组成部分：U-Net、图像编码器（变分自编码器，VAE）和文本编码器（CLIP）。

  • U-Net：Unet 处理带噪潜在空间 (x) 以预测噪声，利用条件模型，该模型还结合了时间步 (t) 和文本嵌入以进行引导。

  • VAE：VAE 模型配备了编码器和解码器，在潜在扩散训练期间进行图像压缩。例如，在标准的 Stable Diffusion 训练阶段，输入图像从 512x512x3 维度压缩到 64x64x4。与像素空间扩散模型相比，这种压缩降低了内存和计算要求。随后，在推理过程中，解码器通过将去噪后的潜在表示转换回其原始、有形的图像形式来逆转此过程。

  • 文本编码器：文本编码器通常是一个像 CLIP 这样的简单 Transformer，它将输入提示转换为嵌入，从而引导 U-Net 的去噪过程。这些嵌入有助于训练 U-Net 有效地处理带噪潜在空间。

模型配置#

在本节中，我们将解释如何配置 Stable Diffusion 模型的 VAE、U-Net 和文本编码器组件的大小和初始化。

变分自编码器#

VAE 配置在 first_stage_config 下定义。

first_stage_config:
    _target_: nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.autoencoder.AutoencoderKL
    from_pretrained: /path/to/vae.bin
    embed_dim: 4
    monitor: val/rec_loss
    ddconfig:
      double_z: true
      z_channels: 4
      resolution: 256  #Never used
      in_channels: 3
      out_ch: 3
      ch: 128
      ch_mult:
      - 1
      - 2
      - 4
      - 4
      num_res_blocks: 2
      attn_resolutions: []
      dropout: 0.0
    lossconfig:
      target: torch.nn.Identity

VAE 权重在训练期间是固定的，并且必须将预训练的检查点传递给 first_stage_config.from_pretrained 以进行初始化。VAE 架构在 Stable diffusion v1 和 v2 系列中是共享的。VAE 的缩放因子为 2**(len(ch_mult - 1))，在本例中为 8。因此，输出图像形状将为 (H//8, W//8, 4)。

U-Net#

U-Net 配置在 unet_config 下定义。

unet_config:
    _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.openaimodel.UNetModel
    from_pretrained: /path/to/pretrain.ckpt
    from_NeMo: True #Must be specified when from pretrained is not None, False means loading unet from HF ckpt
    image_size: 32 # unused
    in_channels: 4
    out_channels: 4
    model_channels: 320
    attention_resolutions:
    - 4
    - 2
    - 1
    num_res_blocks: 2
    channel_mult:
    - 1
    - 2
    - 4
    - 4
    num_head_channels: 64
    use_spatial_transformer: true
    use_linear_in_transformer: true
    transformer_depth: 1
    context_dim: 1024
    use_checkpoint: False
    legacy: False
    use_flash_attention: True

  • 如果未指定 from_pretrained，则 U-Net 使用随机权重初始化。要进行微调，您可以提供预训练的 U-Net 检查点，可以来自中间 NeMo 检查点（设置 from_NeMo=True），也可以来自 Huggingface 等其他平台（设置 from_NeMo=False）。

  • U-Net 大小

    • num_res_blocks：定义每个级别的 resnet 块的计数。

    • model_channels 和 channel_mult：设置每个级别的张量维度。

  • 注意力块

    • attention_resolution：在每个级别的 resnet 块之后集成注意力块。

    • use_spatial_transformer：指定使用的注意力块的类型。

    • use_linear_in_transformer：在线性层和卷积层之间选择用于输入/输出投影。

    • transformer_depth：指示每个 spatial_transformer_block 中 basic_transformer_block 的计数。

  • context_dim：必须调整以匹配文本编码器的输出维度。

文本编码器#

文本编码器配置在 cond_stage_config 下定义。

要在 stable diffusion 中使用 CLIP 模型的 NeMo 实现，可以使用以下 cond_stage_config

cond_stage_config:
  _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.FrozenMegatronCLIPEmbedder
  restore_from_path: /path/to/nemo_clip.nemo
  device: cuda
  freeze: True
  layer: "penultimate"

  • restore_from_path：必须提供以使用 NeMo CLIP 模型，所有 CLIP 配置相关信息都已嵌入在 .nemo 检查点文件中。

  • layer：指定哪个层的输出将用作文本编码器输出。

或者，也可以使用 Huggingface 实现的 CLIP 模型，使用以下配置

cond_stage_config:
    _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.FrozenOpenCLIPEmbedder
    arch: ViT-H-14
    version: laion2b_s32b_b79k
    device: cuda
    max_length: 77
    freeze: True
    layer: "penultimate"

  • arch 和 version：确定要加载哪个 CLIP 模型。

使用预缓存潜在空间进行训练#

由于 VAE 和文本编码器在训练期间保持冻结，您可以离线预先计算图像和标题潜在空间，从而提高训练吞吐量。要创建预缓存数据集，请参阅 多模态数据集。对于使用此数据集进行训练，请正确配置 model.data 部分，并将 model.first_stage_key=image_encoded 与 model.cond_stage_key=captions_encoded 一起设置。

参考#

本页内容