重要提示
您正在查看 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 应该是相应的控制输入。