使用 torch.compile 后端编译 Stable Diffusion 模型

此交互式脚本旨在作为 Torch-TensorRT 工作流与 Stable Diffusion 模型上的 torch.compile 示例。示例如下

导入和模型定义

import torch
import torch_tensorrt
from diffusers import DiffusionPipeline

model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda:0"

# Instantiate Stable Diffusion Pipeline with FP16 weights
pipe = DiffusionPipeline.from_pretrained(
    model_id, revision="fp16", torch_dtype=torch.float16
)
pipe = pipe.to(device)

backend = "torch_tensorrt"

# Optimize the UNet portion with Torch-TensorRT
pipe.unet = torch.compile(
    pipe.unet,
    backend=backend,
    options={
        "truncate_long_and_double": True,
        "enabled_precisions": {torch.float32, torch.float16},
    },
    dynamic=False,
)

推理

prompt = "a majestic castle in the clouds"
image = pipe(prompt).images[0]

image.save("images/majestic_castle.png")
image.show()