AITemplate&StableDiffusion をWindows (WSL)で動かしてみた

Meta社が開発してる、オープンソース高性能推論エンジン「AITemplate」
を使ってStableDiffusionを動かしてみました。
実際かなり早くなったので、興味ある方は是非。導入もそんなに難しくないです。
image2imageもサポートしてくれてるのも嬉しい。
(ただ、RTX30以降やA100とかじゃないと使えないのでそこは注意

WSL上のdockerで動かすまでの手順になります。
WSLのinstallは割愛します。

Dockerのinstall

Docker Desktop for Windowsを入れれば使えるようになるので
下記の公式参考にDockerインストールしましょう。

WSL での Docker コンテナーの概要

NVIDIA Container Toolkitのインストール

dockerのコンテナから GPU 使えるようにするために
NVIDIA Container ToolkitをWSL上でインストールします。
下記の公式通りに手順を踏めば問題なくインストールできるはずです。

Installation Guide — NVIDIA Cloud Native Technologies documentation

AITemplateリポジトリのダウンロード

こちらのリポジトリをもってきてください。

※ --recursive を忘れずにつけること！

git clone --recursive https://github.com/facebookincubator/AITemplate

Dockerイメージのビルド

cd AITemplate
./docker/build.sh cuda

コンテナ起動

docker  run  --gpus all  --name ait -it  ait:latest /bin/bash

# 起動後にdocker内のシェルでGPUの利用が出来るかチェック
nvidia-smi

必要モジュールのインストール

ここからはコンテナ上での作業になります。

Stable Diffusionを動かすのに必要はpythonのモジュールをinstallします。

pip install transformers==4.25.0
pip install diffusers==0.11.0
pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu113
pip install accelerate

サンプルPJへ

cd examples/05_stable_diffusion/

モデルDL

サンプル起動様にモデルをDLしたいので
huggingfaceのtokenを取得してください。

https://huggingface.co/settings/tokens

取得したtokenで下記コマンドを実行

python3 scripts/download_pipeline.py --token 取得したトークン

コンパイル

結構時間かかります。

python3 scripts/compile.py

デモの起動

https://github.com/facebookincubator/AITemplate/tree/main/examples/05_stable_diffusion

python3 scripts/demo.py

# i2iのデモスクリプト
python3 scripts/demo_img2img.py

これで無事動くと思います。
スクリプトも見てもらうと通常のdiffusersを使って書くコード同じように書ける（様にラッパーを用意いしてくれてるので)使い勝手も大分良いと思います。

Waifudiffusionなどの2次元系モデル動かしたい場合はさらにこちらの手順でいけます。