重要提示

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

Flux#

源自 Flux。Flux 是一个基于 Transformer 的 120 亿参数扩散模型,旨在从标题提示生成高质量图像。NeMo 实现的 Flux 使用 Megatron Core 重新实现了所有 Transformer 层,并通过自定义 FSDP 和各种并行性(即将推出 Megatron-LM 的最新更新)进行了分布式训练增强。

从 Hugging Face 导入到 NeMo 2.0#

Hugging Face 检查点可以在 FLUX.1-dev 找到。要在 NeMo 中运行 Flux 的训练或推理管道,请在本地下载检查点并提供模型配置的路径。它将自动转换为匹配 NeMo 结构。

训练#

对于训练,请在预设配方中将检查点路径提供给模型配置,如下所示

@run.cli.factory(target=llm.train)
def convergence_test() -> run.Partial:
    recipe = flux_training()
    recipe.model.flux_params.flux_config = run.Config(
        FluxConfig,
        ckpt_path='/path/to/FLUX.1-dev/transformers',
        do_convert_from_hf=True,
        save_converted_model_to='/ckpts'
    )
    recipe.model.flux_params.t5_params = run.Config(
        T5Config, version='/ckpts/text_encoder_2'
    )
    recipe.model.flux_params.clip_params = run.Config(
        ClipConfig, version='/ckpts/text_encoder'
    )
    recipe.model.flux_params.vae_config = run.Config(
        AutoEncoderConfig, ckpt='/ckpts/ae.safetensors', ch_mult=[1, 2, 4, 4], attn_resolutions=[]
    )
    return recipe

上面的命令将转换后的文件保存在 save_converted_model_to 中指定的路径中,并将其命名为 nemo_flux_transformer.safetensors。如果未提供 `save_converted_model_to`,则仍将执行转换,但不会保存转换后的检查点。

在后续运行中,当向 flux_config 提供 nemo_flux_transformer.safetensors 的路径时,应将 do_convert_from_hf 设置为 False。

VAE(变分自动编码器)、CLIP 和 T5 等其他组件的检查点在配方中以类似方式设置,如上面的示例所示。

推理#

对于 Flux 推理,提供的脚本也接受 Hugging Face 检查点,类似于训练脚本。

参数说明#

  • flux_ckpt Flux Transformer 检查点的路径。接受的格式包括从上面链接下载的 Transformer 文件夹中的 safetensors(例如,/FLUX-1.dev/transfomer/

或 NeMo 格式的转换后的 safetensors。

  • do_convert_from_hf:要将下载的检查点转换为 NeMo 格式,需要 –do_convert_from_hf 作为

命令行参数。

  • save_converted_model_to:转换后的 Flux Transformer 权重将默认本地保存到 /ckpts

用户可以通过 –save_converted_model_to 指定此路径。

  • vae_ckpt:VAE 检查点的路径,可以在上面链接的 FLUX-1.dev/ae.safetensors/ 中找到。

  • clip_version:用于处理输入提示的 CLIP 文本编码器版本,例如 openai/clip-vit-large-patch14,它

将从服务器自动下载。或者,您可以直接提供本地存储的 CLIP 检查点的路径,例如,FLUX-1.dev/text_encoder

  • t5_version:用于处理输入提示的 T5 文本编码器版本,例如 google/t5-v1_1-xxl,它

将从服务器自动下载。或者,您可以直接提供本地存储的 T5 检查点的路径,例如,FLUX-1.dev/text_encoder_2

  • width & height:最终输出图像的尺寸。

  • inference_steps:每个提示运行的去噪步骤数。

  • num_iamges_per_prompt:每个提示生成的推理图像数量。

  • prompts:用于图像生成的文本提示。当在单个运行中提供多个提示时,使用 , 分隔

不同的提示。

  • guidance:引导比例控制提示条件的强度。

脚本的使用示例#

torchrun flux_infer.py --flux_ckpt /ckpts/nemo_flux_transformer.safetensors \
--prompts  'A cat holding a sign that says hello world'
--inference_steps 50 --guidance 3.5

训练配方#

我们为 Flux 提供了预定义的训练配方,可用于性能基准测试和使用 NeMo-Run 的收敛性测试。这些配方托管在 flux_training.py 中。

默认情况下,Flux 模型将被随机初始化,Flux Transformer 层数和 Flux 单个 Transformer 层数均设置为 1。此默认配方还使用 MockDataModule,它为图像和示例文本输入生成随机张量。

为了使用您自己的数据集启动训练,请按照以下示例使用您自己的配置覆盖默认配方

@run.cli.factory(target=llm.train)
def convergence_test() -> run.Partial:
    recipe = flux_training()  ##Predefined base recipe where most default values are set
    ### Load checkpoints for embedder components
    recipe.model.flux_params.t5_params = run.Config(
        T5Config, version='/ckpts/text_encoder_2'
    )
    recipe.model.flux_params.clip_params = run.Config(
        ClipConfig, version='/ckpts/text_encoder'
    )
    recipe.model.flux_params.vae_config = run.Config(
        AutoEncoderConfig, ckpt='/ckpts/ae.safetensors', ch_mult=[1, 2, 4, 4], attn_resolutions=[]
    )
    recipe.model.flux_params.device = 'cuda'
    recipe.trainer.devices = 8         ### Number of GPUs per node
    ### Change data module from MockDataModule to DiffusionDataModule
    ### Provide the path to Megatron-Energon Compatible dataset
    recipe.data = flux_datamodule('/dataset/fill50k/fill50k_tarfiles/')
    recipe.model.flux_params.flux_config = run.Config(
        FluxConfig,
        num_joint_layers=19,
        num_single_layers=38
    )   ## Adjust the number of transfomer layers
    return recipe

这是一个运行您自己的自定义配方的示例命令

脚本的使用示例

torchrun --nproc_per_node 8 flux_training.py --yes --factory convergence_test

Flux ControlNet#

训练#

Flux ControlNet 训练脚本 flux_controlnet_training.py 是使用 NeMo-Run 实现的。FluxControlNetConfig 包括初始化 Flux 模型所需的所有组件,以及确定 ControlNet 部分架构的配置。例如,

@run.cli.factory(target=llm.train)
def unit_test() -> run.Partial:
    ...
    recipe.model.flux_controlnet_config.num_single_layers = 1
    recipe.model.flux_controlnet_config.num_joint_layers = 1
    # Initialize transformer blocks in ContolNet from the weights of corresponding flux blocks
    recipe.model.flux_controlnet_config.load_from_flux_transformer = True

推理#

Flux ControlNet 推理脚本位于 flux_controlnet_infer.py,其操作方式与标准 Flux 推理脚本类似。但是,除了 Flux 所需的检查点外,Flux ControlNet 还需要额外的输入。

其他参数#

  • controlnet_ckpt:此参数接受来自 Hugging Face 的 ControlNet 检查点。当使用 –do_convert_from_hf 标志时,Flux 检查点和

Flux ControlNet 检查点都将被转换并保存以用于 Flux 推理。

  • control_image:用作推理管道中附加条件的控制图像的路径。

  • num_joint_layers:此参数现在确定 ControlNet 中 Transformer 层的数量。Flux 模型中 Transformer 层的数量保持与 FLUX.1-dev 中相同的默认设置。num_single_layers 也适用。

Energon 兼容数据集#

有关使用 Energon 准备数据集的详细信息,请参阅 数据准备 部分。对于 Flux 用法,我们采用了 WebDataset 格式,其中图像-文本对保存在具有相同名称但不同扩展名的 tar 文件中。

.
├── 001.txt
├── 001.jpg
├── 002.txt
└── 002.jpg

flux_datamodule 应由 Energon 准备为 CrudeDataset。我们使用虚拟任务编码器来加载原始图像和文本输入。基本数据增强将很快添加到 flux_datamodule 中。

对于训练 Flux ControlNet,.jpg 文件用作原始图像,而 .png 应该是相应的控制输入。