Japanese Stable CLIP を試してみた

Stable Diffusion (WebUI) を TensorRT で高速化


TensorRT を使うことで、Stable Diffision の推論を高速化することができるようです。
TensorRT の導入がいろいろ手順があって大変なようですが、今回はその手順をコンテナを使って楽して試してみましょう。
なお、VRAM 8GB の GeForce 2070 でも実行可能でした。

前提

・RTX 世代の GeForce
・Docker で GPU が使えるようになっている

事前準備

・GeForce のドライバを更新する
(多分) 2023年10月にリリースされている最新版のドライバをインストールする必要があります。下記の TensorRT の拡張機能の GitHub には、「Windows: >=452.39」という記載がありますので、それ以上のバージョンにすれば OK だと思います。

・お好きな Stable Diffision のモデルをダウンロードしておく
※SDXL は、2023年10月時点では対応していない様子です。
⇒拡張子 .safetensors を /work ディレクトリに保存しました。

・NGC の TensorRT コンテナイメージをダウンロードする
docker pull nvcr.io/nvidia/tensorrt:23.09-py3

・TensorRT の拡張機能をダウンロードする
⇒今回は ZIP でダウンロードして、/work ディレクトリ配下に保存しました。

実行

・WebUI 用のポートと、モデルデータ共有用のディレクトリを設定してコンテナを起動する
docker run -it --gpus=all --rm -v /work:/work -p 7860:7860 nvcr.io/nvidia/tensorrt:23.09-py3 /bin/bash

・必要なパッケージをインストールする
apt update
apt install -y python3.10-venv libgl1-mesa-dev git vim libglib2.0-0 curl

・Stable Diffision WebUI をダウンロードする
git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
cd stable-diffusion-webui/
sed -i -e s/can_run_as_root=0/can_run_as_root=1/ webui.sh
※root でも起動できるように sed コマンドを実行します

・モデルをコピーする
cp /work/xxx.safetensors model/Stable-diffusion

・TensorRT 拡張機能を配置する
pushd extensions
unzip /work/Stable-Diffusion-WebUI-TensorRT-main.zip
popd

・WebUI を起動する
COMMANDLINE_ARGS="--listen --enable-insecure-extension-access" ./webui.sh
※初回は必要なライブラリをインストールするので、10分程度かかると思います

・WebUI が起動したら、ブラウザから http://localhost:7860 を開く
「TensorRT」タブがあるはずなので、使用するモデルをプルダウンから選んでおき、「Export Default Engine」をクリックする
※10分程度処理に時間がかかります。ターミナルに処理状況が表示されます

タブの下部に「Export Successfully」と表示されたら、ターミナルで Ctrl+C を押し、いったん WebUI を終了する。

・再度 WebUI を起動し、ブラウザから http://localhost:7860 を開く
TensorRT タブの一番下、Available TensorRT Engine Profile を開くと、利用可能な解像度などが表示されます。この設定にすると、TensorRT が効いて推論が速くなります。


TensorRT を使わないと 1 枚 3,4 秒かかるのが、1 枚 1 秒かかるかどうか、といった程度まで短縮できます。画像生成がはかどりますね。

コメント