注意
跳转到末尾 下载完整示例代码
使用 Torch-TensorRT dynamo 后端编译 FLUX.1-dev 模型¶
本示例展示了如何使用 Torch-TensorRT 优化最新的模型 FLUX.1-dev。
FLUX.1 [dev] 是一个拥有 120 亿参数的 rectified flow Transformer,能够从文本描述生成图像。这是一个用于非商业应用的开放权重、指导蒸馏模型。
要运行此演示,您需要获得 Flux 模型访问权限(如果尚未获得,请在 FLUX.1-dev 页面申请)并安装以下依赖项
pip install sentencepiece=="0.2.0" transformers=="4.48.2" accelerate=="1.3.0" diffusers=="0.32.2" protobuf=="5.29.3"
FLUX.1-dev
流水线包含不同的组件,例如 transformer
、vae
、text_encoder
、tokenizer
和 scheduler
。在本示例中,我们演示了如何优化模型的 transformer
组件(该组件通常占整个端到端扩散延迟的 >95%)
导入以下库¶
import torch
import torch_tensorrt
from diffusers import FluxPipeline
from torch.export._trace import _export
定义 FLUX-1.dev 模型¶
使用 FluxPipeline
类加载 FLUX-1.dev
预训练流水线。FluxPipeline
包含生成图像所需的各种组件,例如 transformer
、vae
、text_encoder
、tokenizer
和 scheduler
。我们使用 torch_dtype
参数以 FP16
精度加载权重
DEVICE = "cuda:0"
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
torch_dtype=torch.float16,
)
# Store the config and transformer backbone
config = pipe.transformer.config
backbone = pipe.transformer.to(DEVICE)
使用 torch.export 导出主干网络¶
定义 dummy 输入及其对应的动态形状。由于 0/1 特化,我们以 batch_size=2
的动态形状导出 transformer 主干网络
batch_size = 2
BATCH = torch.export.Dim("batch", min=1, max=2)
SEQ_LEN = torch.export.Dim("seq_len", min=1, max=512)
# This particular min, max values for img_id input are recommended by torch dynamo during the export of the model.
# To see this recommendation, you can try exporting using min=1, max=4096
IMG_ID = torch.export.Dim("img_id", min=3586, max=4096)
dynamic_shapes = {
"hidden_states": {0: BATCH},
"encoder_hidden_states": {0: BATCH, 1: SEQ_LEN},
"pooled_projections": {0: BATCH},
"timestep": {0: BATCH},
"txt_ids": {0: SEQ_LEN},
"img_ids": {0: IMG_ID},
"guidance": {0: BATCH},
"joint_attention_kwargs": {},
"return_dict": None,
}
# The guidance factor is of type torch.float32
dummy_inputs = {
"hidden_states": torch.randn((batch_size, 4096, 64), dtype=torch.float16).to(
DEVICE
),
"encoder_hidden_states": torch.randn(
(batch_size, 512, 4096), dtype=torch.float16
).to(DEVICE),
"pooled_projections": torch.randn((batch_size, 768), dtype=torch.float16).to(
DEVICE
),
"timestep": torch.tensor([1.0, 1.0], dtype=torch.float16).to(DEVICE),
"txt_ids": torch.randn((512, 3), dtype=torch.float16).to(DEVICE),
"img_ids": torch.randn((4096, 3), dtype=torch.float16).to(DEVICE),
"guidance": torch.tensor([1.0, 1.0], dtype=torch.float32).to(DEVICE),
"joint_attention_kwargs": {},
"return_dict": False,
}
# This will create an exported program which is going to be compiled with Torch-TensorRT
ep = _export(
backbone,
args=(),
kwargs=dummy_inputs,
dynamic_shapes=dynamic_shapes,
strict=False,
allow_complex_guards_as_runtime_asserts=True,
)
Torch-TensorRT 编译¶
注意
编译需要具有大内存(> 80GB)的 GPU,因为 TensorRT 以 FP32 精度存储权重。这是一个已知问题,将在未来版本中解决。
我们通过设置 use_fp32_acc=True
启用 FP32
矩阵乘法累加,通过引入转换为 FP32
的节点来确保精度得以保留。我们还启用了显式类型指定,以确保 TensorRT 遵守用户设置的数据类型,这是 FP32 矩阵乘法累加的要求。由于这是一个 120 亿参数的模型,在 H100 GPU 上编译大约需要 20-30 分钟。该模型完全可转换,并生成一个单一的 TensorRT 引擎。
trt_gm = torch_tensorrt.dynamo.compile(
ep,
inputs=dummy_inputs,
enabled_precisions={torch.float32},
truncate_double=True,
min_block_size=1,
use_fp32_acc=True,
use_explicit_typing=True,
)
后处理¶
释放由导出程序和 pipe.transformer 占用的 GPU 内存 将 Flux 流水线中的 transformer 设置为 Torch-TRT 编译后的模型
del ep
backbone.to("cpu")
pipe.to(DEVICE)
torch.cuda.empty_cache()
pipe.transformer = trt_gm
pipe.transformer.config = config
使用提示词生成图像¶
提供提示词和要生成图像的文件名。这里我们使用提示词 A golden retriever holding a sign to code
。
# Function which generates images from the flux pipeline
def generate_image(pipe, prompt, image_name):
seed = 42
image = pipe(
prompt,
output_type="pil",
num_inference_steps=20,
generator=torch.Generator("cuda").manual_seed(seed),
).images[0]
image.save(f"{image_name}.png")
print(f"Image generated using {image_name} model saved as {image_name}.png")
generate_image(pipe, ["A golden retriever holding a sign to code"], "dog_code")
生成的图像如下所示

脚本总运行时间:( 0 分钟 0.000 秒)