Radeon RX 6800でStable Diffusionを動かしてみたメモ

最近話題の画像生成AI、Stable Diffusionを手元のローカル環境のPCのAMD Radeon RX 6800で動かしてみました。慣れない設定が多くて難儀しましたが、なんとか動かす事ができたので、ここにメモを残します。業界標準のGPGPUプラットフォームであるNVIDIA CUDAではなく、そのAMD版と言えるROCmを使う事になるため、Windowsでそのまま動くわけではなかったりとやや複雑ですが、一度設定してしまえば利用料金や回数制限などを気にする事なく使い放題です。

なお、この手順はYouTubeで公開されていた動画 Stable Diffusion (DALLE-2 clone) on AMD GPU を参考にしていますが、試行錯誤の結果、動画とは少し違う手順となっています。なお、筆者は機械学習素人故、おかしな記述もあるかと思います。その場合はご指摘いただけましたら幸いです。

本記事は以下の環境を前提としています。

  • PC/AT互換機(いわゆるMacやRaspberry Pi等ではないPC)
  • AMD Radeon GPU
  • 8GB以上程度のUSBメモリ
  • LinuxをインストールするためのHDD, SSD等ストレージデバイス(256GB以上程度を推奨)
  • インターネット接続

Apple Silicon Macの場合、NVIDIA GeForce GPUの場合はそれぞれ別の方が書かれた記事があるのでそちらを参照してください。

おおまかな手順は次の通りです。Stable Diffusionが動作する環境を整備するため、コンテナ仮想化プラットフォーム Docker 、およびPython用パッケージ管理システム condapip を使用します。

Linuxディストリビューションを実環境にインストール

Stable DiffusionをGPUで使うにはNVIDIA CUDAのAMD版と言えるROCmが必要となりますが、ROCmにはWindows版が存在しません。このため、使用にはLinux OSの動作する環境が必要となります。なお、Virtualbox等のバーチャルマシンVM環境ではGPUの全機能がホストOSからゲストOS側に露出していない事があるため、VMではなく実環境が必要となります。(最近のバージョンのWSL2ではGPUを叩けるらしいという話も一部目にしましたが、未検証です)

僕は Ubuntu Linux 22.04 LTS 日本語 Remix をダウンロードしました。実際のところ、2022年8月現在のバージョンのROCmで公式サポートされているのはUbuntu 20.04 LTSなのですが、こちらで試したところ有線LANが動かないなど無関係の問題が色々発生したため、面倒なので22.04で進める事にしました。

Linuxディストリビューションのインストールに関しては他で詳しく説明されているページが多数あるため細かくは省きますが、

  • 現在のOSで、 Ubuntu Linux 22.04 LTS 日本語 Remix ダウンロードページからISOイメージファイルをダウンロード
  • USBメモリやSDカードにISOイメージファイルを焼き込むツール BalenaEtcher をダウンロード、インストール
  • PCにUSBメモリを刺し、BalenaEtcherを使ってUbuntuのISOイメージをUSBメモリに焼き込む
  • USBメモリを刺してPCを再起動し、起動時にDeleteキー等(マザーボードによって異なります)を押してUSBメモリから起動
  • 起動後、「Ubuntuをインストール」を選び、画面上の手順に沿って用意しておいたストレージデバイスにUbuntu Linuxをインストール

ROCmをインストール

インストールが終わったら、インストールされたLinux環境で起動し、ターミナルを開き、ROCmをインストールします。RadeonOpenCompute GitHubのこちらのページの手順を参考に一部改変しています。なお、sudo (SuperUser Do)から始まる行は管理者権限で実行するため、パスワードを聞かれる場合があります。

# ダウンロード用ツールやバージョン管理ツールが入ってなければインストール(後々使います)
sudo apt install wget curl git
# ROCm用カーネルモジュールを導入
wget -qO - http://repo.radeon.com/rocm/rocm.gpg.key | sudo apt-key add -
# 元ページの記述ではレポジトリのURLが古かったようなので以下に変更
echo deb [arch=amd64] https://repo.radeon.com/rocm/apt/latest/ ubuntu main | sudo tee /etc/apt/sources.list.d/rocm.list

(必要であれば)ROCmインストールをUbuntu 22.04に対応させる

この後 sudo apt-get update && sudo apt-get install rock-dkms と打ち込めばROCmカーネルモジュールがインストールされる…はずですが、Ubuntu 22.04はROCmに公式サポートされていないためそのままだとエラーが出ます。その回避方法はこちらのAsk Ubuntu Q&Aページに載っていました。

15ステップ+αもあり複雑ですが、順番に従っていったところ無事インストールできました。(Ubuntu 20.04ではこの手順は不要なはずです。また、いずれROCmがUbuntu 22.04を公式サポートしたらこの手順は不要となるはずです。)

ROCmをインストール(続)

気を取り直して手順を続けます。

sudo apt-get update && sudo apt-get install rock-dkms
sudo update-initramfs -u
# 一旦再起動
sudo reboot
#再起動後、再びターミナルを開き、現在のユーザーをvideoグループとrenderグループに追加
sudo usermod -a -G video $LOGNAME
sudo usermod -a -G render $LOGNAME

Dockerをインストール

コンテナ仮想化プラットフォームDockerをインストールします。

curl -sSL https://get.docker.com/ | sh

インストール完了後、以下のコマンドを入力してDocker環境の情報を表示します。

sudo docker info

ここで Storage Driveroverlay2 になっていればOKです。僕は最初からなっていたので特に何もせず進めました。 devicemapper になっていた場合、イメージやコンテナのサイズが最大10GBに制限されるため overlay2 に切り替えるべきとのことです。

所有するグラフィックカードに合致するDockerイメージを選ぶ

Dockerイメージの公開サイトDocker Hubで、AMDによりROCm対応PyTorchのイメージが公開されています。下準備として、まずは概要欄に書かれている下記ショートカットコマンドをコピーし、自分のホームディレクトリの ~/.bashrc ファイルの最後尾に追記します:

# GPUなどを適切にホストOS側から露出させ、またdockerxディレクトリをホストOS側からマウントした状態でコンテナを実行できるショートカット
alias drun='sudo docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v $HOME/dockerx:/dockerx'
# ついでにこちらも.bashrcに追記しておくと便利
# 最後に作ったコンテナを起動し直して中に入るショートカット
alias redock='sudo docker start -i $(sudo docker ps -q -l)'

前述のROCm対応PyTorchイメージのページの「Tags」タブを開くと、さまざまなAMD GPU用のDockerイメージが一覧となっています。ここで自分の持っているグラフィックカードのアーキテクチャに合わせたイメージを選び、右側の書類マークをクリックして名前をクリップボードにコピーします。僕の場合はRadeon RX 6800(Navi 21アーキテクチャ)のため、 rocm/pytorch:rocm5.2_ubuntu20.04_py3.7_pytorch_1.11.0_navi21 を選びました。わからない場合は latest で動くかもしれません。なお、Radeon RX 6700 XTなど、機種によっては認識させるために環境変数を追記する必要があるとのことです(未確認)。

Dockerfileおよびenvironment.yamlをダウンロード

本記事が参考とするチュートリアル動画ではこの後イメージからコンテナを立ち上げた後、コンテナ内で手作業でconda環境の構築やROCmに合わせた改変などを行っていましたが、これを自動化するためDockerfileとconda設定用のenvironment.yaml改変版を作りました。こちらのGitHubレポジトリからダウンロードし、以下のようにコンテナをビルドして起動します。なお、ここでダウンロードされるDockerfileでは出発元のイメージを前述のRX6800用のrocm/pytorch:rocm5.2_ubuntu20.04_py3.7_pytorch_1.11.0_navi21 にしてありますが、違うグラフィックカードをお持ちの方は適宜最初の行のFROM文を適切なイメージ名に書き換えてください。

git clone https://github.com/heistak/stable-diffusion-radeon.git
sudo docker build --tag stable-diffusion-radeon:1.0 stable-diffusion-radeon
# Dockerfileの中ではStable Diffusionレポジトリを取得し、
# conda環境設定用の改変版yaml(environment2.yaml)をそこに追加、
# またその環境設定ファイルに基づいてROCm対応版pyTorch等をセットアップします
# しばらく時間がかかります
# 構築が完了したのちに前述のショートカットエイリアスを使ってコンテナを起動
drun stable-diffusion-radeon:1.0

Stable Diffusionレポジトリをコピー

drunエイリアスコマンドを使ってDockerコンテナが立ち上がった際、ホームディレクトリ下に~/dockerx/というディレクトリが出来ています。このディレクトリはホストOS側のファイルシステムに所属しますが、コンテナの動作中はコンテナ側からも見えており読み書きが可能です。コンテナ側からはルートディレクトリ直下の/dockerxの位置に見えています。

コンテナ内でここに移動し、rocmというディレクトリを作り、その下にStable Diffusion本体のレポジトリをクローンする…ところですが、実はStable Diffusionのレポジトリは前項のDockerコンテナのビルド時の自動化された流れの中で既に別の場所(/root/)にクローン済みなので、それをコピーして持ってきます。(移動だとうまくいきません)

mkdir /dockerx/rocm
cp -r /root/stable-diffusion /dockerx/rocm/stable-diffusion

Weightsをダウンロード

ターミナルの外でブラウザを開き、機械学習コミュニティサイトHugging Face内のページからWeightsファイルをダウンロードします。ページ右上のSign Upボタンでユーザー登録をした後、Access Repositoryのボタンを押すとダウンロードできるようになります。ダウンロードするのは sd-v1-4.ckpt の方で大丈夫です。

「ファイル」アプリ、あるいはもう一個別のシェルなどを使って、~/dockerx/rocm/stable-diffusion/models/ldm/ 下に stable-diffusion-v1 というディレクトリを作り、その中にダウンロードしたファイルを model.ckpt という名前にして配置します。

# これはDockerコンテナ内ではなくコンテナ外のホストOS側で開いたシェルで操作する
sudo mkdir ~/dockerx/rocm/stable-diffusion/models/ldm/stable-diffusion-v1
cd ~/ダウンロード/ # ダウンロードフォルダの名前は環境によって異なるので適宜読替え
sudo mv ./sd-v1-4.ckpt ~/dockerx/rocm/stable-diffusion/models/ldm/stable-diffusion-v1/model.ckpt

実行

準備が整ったはずなので、コンテナ内のシェルで以下を実行してください。

cd /dockerx/rocm/stable-diffusion
python3 scripts/txt2img.py --prompt "A beautiful sunset on the ocean horizon" --H 512 --W 512 --n_iter 1 --ddim_steps 50 --n_samples 1

初回実行時は色々ダウンロードするので時間がかかります。
“Enjoy”の表示が出て実行が終わった後、~/dockerx/rocm/stable-diffusion/outputs/txt2img-samples/ の中に、水平線に沈む夕日の画像が出ていれば成功です。

なお、コンテナの中ではrootユーザーで動かしている関係上、outputディレクトリ内で生成される画像のオーナーがroot:rootになりますが、適宜ホストOS側で sudo chown -R $USER:$USER ~/dockerx/rocm/stable-diffusion/outputs/* 等で自分の所有に直してください。

2回目以降の使用

2回目以降の使用時には、前に作ったコンテナの名前をdocker ps -aで調べ、docker start -i <コンテナ名>で再度立ち上げれば使えます。途中の手順で~/.bashrcファイルにエイリアス行を入れてあれば、 redock と打つだけで最後に作ったコンテナに入り直せるはずです。