使用 torch.compile 後端編譯 Stable Diffusion 模型

此互動式腳本旨在作為 Stable Diffusion 模型上使用 torch.compile 的 Torch-TensorRT 工作流程範例。以下為範例輸出

匯入和模型定義

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()