Stable Diffusionを高速化できる裏ワザ

torch_tensorrtでGPU最適化を行う

以下のように、torch_tensorrtにコンパイルすることで速くなる。

T4 GPU

	sec
torch	8.1
torch_tensorrt	6.1

A100 GPU

	sec
torch	2.1
torch_tensorrt	1.1

使い方も簡単。

使い方

インストール。

pip install diffusers
pip install torch_tensorrt

diffusionモデルの初期化

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

実行。

import time
prompt = "robot, penguin"
seed = 42
generator = torch.manual_seed(seed)

image = pipe(prompt).images[0]
image.save("robot_penguin.png")
image.show()

ただし、仕上がりが少しだけ異なる。

torch

torch_tensorrt

🐣

---

フリーランスエンジニアです。
AIについて色々記事を書いていますのでよかったらプロフィールを見てみてください。

もし以下のようなご要望をお持ちでしたらお気軽にご相談ください。
AIサービスを開発したい、ビジネスにAIを組み込んで効率化したい、AIを使ったスマホアプリを開発したい、
ARを使ったアプリケーションを作りたい、スマホアプリを作りたいけどどこに相談したらいいかわからない…

いずれも中間コストを省いたリーズナブルな価格でお請けできます。

お仕事のご相談はこちらまで
[e-mail]

機械学習やAR技術を使ったアプリケーションを作っています。
機械学習／AR関連の情報を発信しています。

X
Medium
GitHub