重要提示

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

ControlNet#

模型介绍#

ControlNet [MM-MODELS-CN1] 是一种神经网络结构，通过添加额外的条件来控制扩散模型。它将神经网络块的权重复制到“锁定”副本和“可训练”副本中。“可训练”副本学习您的条件。“锁定”副本保留您的模型。通过这种方式，ControlNet 可以重用 SD 编码器作为深度、强大、稳健和强大的骨干，以学习各种控制。NeMo Multimodal 提供了基于分割图生成图像的训练管道和示例实现。用户可以灵活地使用自己的控制输入数据集和配方探索其他实现方式。

ControlNet structure on stable diffusion (See :cite:`mm-models-cn-controlnetgithub`)

ControlNet 数据集#

ControlNet 采用 WebDataset 格式进行数据摄取。（请参阅 数据集）除了以 tarfile 格式保存的基本图像-文本对，文件名相同但扩展名不同（如 000001.jpg 和 000001.txt）之外，ControlNet 还要求 tarfile 中包含控制输入，可通过其特定扩展名识别。默认情况下，控制输入应存储为 000001.png，以便在 NeMo 的实现中正确加载和识别。

模型配置#

即使 Stable Diffusion 权重的原始副本被锁定，也需要正确的配置设置以及兼容的预训练检查点才能进行初始化。有关 unet_config、first_stage_config 和 cond_stage_config 的更多详细信息，请参阅 模型配置。

控制阶段配置#

control_stage_config:
    _target_: nemo.collections.multimodal.models.controlnet.controlnet.ControlNet
    params:
      from_pretrained_unet: /ckpts/v1-5-pruned.ckpt
      from_NeMo: False
      image_size: 32 # unused
      in_channels: 4
      hint_channels: 3
      model_channels: 320
      attention_resolutions: [ 4, 2, 1 ]
      num_res_blocks: 2
      channel_mult: [ 1, 2, 4, 4 ]
      num_heads: 8
      use_spatial_transformer: True
      use_linear_in_transformer: False
      transformer_depth: 1
      context_dim: 768
      use_checkpoint: False
      legacy: False
      use_flash_attention: True

  • from_pretrained_unet：与 unet_config.from_pretrained 相同的逻辑，根据检查点的来源（来自 Huggingface 还是 NeMo）调整 from_NeMo。

  • control_stage_config：概述了 U-Net 可训练副本的架构。所有参数都必须与本节中指定的 U-Net 检查点对齐，这一点至关重要。

  • hint_channels：表示输入控制的通道，在提到的示例中为 3，因为 RGB 图像输入的形状为 (H, W, 3)。

ControlNet 训练选项#

model:
    control_key: hint
    only_mid_control: False
    sd_locked: True
    ...

  • contorl_key：控制输入的标识符，.png 文件将转换为字典，供数据加载器使用，其键为 hint。

  • only_mid_control：设置为 True 时，在训练期间，只有来自可训练副本中间块的输出将合并到锁定副本中。

  • sd_locked：是否在训练期间锁定原始 Stable Diffusion 权重。

参考#

[MM-MODELS-CN1]

Zhang Lllyasviel. Lllyasviel/controlnet. 2023. URL: lllyasviel/ControlNet。

在本页中