キリウ君が読まないノート WSL2のUbuntu20.04でStable Diffusionを動かす (with RTX2060)
オープンソースの画像生成AI
参考にさせていただいた記事
話題のStable Diffusionがオープンソース化されたのでローカルで動かしてみる
WSL2 上の PyTorch に GPU を認識させて深層学習環境をつくる
Ubuntu 20.04へのCUDAインストール方法 - Qiita
筆者環境
- OS … Windows 11 Home Insider Preview
- GPU … GeForce RTX 2060 (VRAM 6GB)
- WSL2のLinuxディストリビューション … Ubuntu 20.04 LTS
- WSL2のメモリ制限 … 16GB
WSLのメモリ制限というのは .wslconfig で指定できる memory のことです。この記事の方法でOptimized Stable Diffusionを実行した時、タスクマネージャで確認したメモリ使用量は最大13GB程度だった。メモリをそれ以下に制限している人は動かないかも。
Windows 10の場合は、然るべきバージョン以上じゃないとWSL2上でGPU計算を行えないらしい。が、筆者はそのあたりでゴチャゴチャやってるうちに、間違えてWindows 11にバージョンアップしてしまった。なのでWindows 10の人は「windows10 wsl2 gpu」とかで調べた方がいいかも。
概要
この記事に記載しているコマンドは、基本的に全てWSL2のUbuntu上で実行すること。
ざっくりまとめると以下の通り。
- WSL2のUbuntuにpyenvをインストールし、pyenvでAnacondaをインストールする。
- WSL2からGPUを使って計算を行うために、Windows側にNVIDIAドライバーをインストールする。
- WSL2からGPUを使って計算を行うために、WSL2側にCUDA Toolkit、CUDA Deep Neural Network (cuDNN) をインストールする。
- WSL2にStable Diffusion本体とStable Diffusionの学習モデルを用意する。
- WSL2でStable Diffusionを実行する。
ただし今回はVRAMが少ないためにStable Diffusion実行時にエラーが出てしまったので、代わりにOptimized Stable Diffusionを使用して継続。とりあえず本来のやり方でやってみて、ダメなら最適化版で動かせればいいという人なら、この記事の手順でできると思う。
それとネットでStable Diffusionの実行方法を調べてるとPyTorchがよく出てくるが、PyTorchは単体でインストールしなくてもStable DiffusionのAnacondaの仮想環境に含まれている。なのでここではインストールしない。
あと、Hugging Face上で公開されてるStable Diffusionの学習モデルにはDiffusers版とオリジナル版がある？みたい。Diffusers版の方が本体と学習モデルを別々に準備しなくてよくて簡単っぽいが、よくわからないのでここでは本体＋オリジナル版で動かす。
それでは次項から。
手順1. pyenvインストール
普通にPythonをインストールしてしまうと後で別のバージョンを切り替えたくなった時に困るので、ここではpyenvを使用する。
以下コマンドでpyenvをインストール → パスを .bashrc に追記 → シェルを再起動。
git clone https://github.com/pyenv/pyenv.git ~/.pyenv
echo 'export PYENV_ROOT="$HOME/.pyenv"' >> ~/.bashrc
echo 'command -v pyenv >/dev/null || export PATH="$PYENV_ROOT/bin:$PATH"' >> ~/.bashrc
echo 'eval "$(pyenv init -)"' >> ~/.bashrc
exec "$SHELL"インストールできたことを以下コマンドで確認。バージョンが表示されればOK。
pyenv -vpyenv 2.3.3-5-ga8afc611手順2. pyenvでAnacondaインストール
Stable DiffusionのREADMEに書いてある手順に conda コマンドが出てくるので、Anacondaをインストールする。
まず以下コマンドで、pyenvでインストール可能なAnacondaのバージョン一覧を確認する。
pyenv install -l | grep anacondaここでは2022年8月26日現時点で最新の anaconda3-2022.05 をインストールする。以下コマンドでインストールし、pyenvが呼び出すPythonのバージョンを anaconda3-2022.05 に切り替える。
pyenv install anaconda3-2022.05
pyenv global anaconda3-2022.05完了したら、以下コマンドで python と conda のバージョンを確認してみる。
python -V
conda -VPython 3.9.12
conda 4.12.0なお、pyenvでインストール済みのPythonのバージョン一覧は以下コマンドで確認可。
pyenv versions system (set by /home/napoporitataso/.pyenv/version)
* anaconda3-2022.05手順3. NVIDIAドライバーインストール
WSL2からGPUを認識できるように、Windows側に『NVIDIA CUDA on WSL driver』をインストールする。これについての公式の説明は以下。
GPU in Windows Subsystem for Linux (WSL) | NVIDIA Developer
上記サイトに記載の通り、『NVIDIA CUDA on WSL driver』はGeForceとQuadroのドライバーに含まれているので、公式サイトから手元のグラボ用のドライバーをダウンロードしてインストールすればよい。公式のダウンロードページは以下。
選択肢は手元の製品や環境に合わせて設定する。筆者の場合は以下の通り。
- Product Type … GeForce
- Product Series … GeForce RTX 20 Series (Notebooks)
- Product … GeForce RTX 2060
- Operating System … Windows 11
- Download Type … Studio Driver (SD)
- Language … Japanese
ダウンロードしたら、普通にWindows上で実行してインストールする。全て初期設定でよい。
完了したらWSL上で以下コマンドを実行して、WSL側からGPUが認識できているか確認する。
nvidia-smiSat Aug 27 01:48:53 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.65.01 Driver Version: 516.94 CUDA Version: 11.7 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 NVIDIA GeForce ... On | 00000000:01:00.0 On | N/A |
| N/A 45C P8 6W / N/A | 74MiB / 6144MiB | 1% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| No running processes found |
+-----------------------------------------------------------------------------+真ん中あたりに『NVIDIA GeForce ...』と表示されており、『6144MiB』とVRAMが6GBであることも表示されているのでOK。
手順4. CUDA Toolkitインストール
WSLにCUDA Toolkitをインストールする。GPUを画像処理ではなく普通の計算のために使用するにはこれが必要らしい。
以下の公式サイトから導入。
CUDA Toolkit 11.7 Update 1 Downloads | NVIDIA Developer
OSやアーキテクチャを選択していくと、自動でインストールコマンドを生成してサイト上に表示してくれる。筆者の場合は以下の通り。
- Operating System … Linux
- Architecture … x86_64
- Distribution … WSL-Ubuntu
- Version … 20.04
- Installer Type … deb (local)
WSLのUbuntuの場合はDistributionに『WSL-Ubuntu』を選択すること。
上記の選択肢の場合、生成されたコマンドは以下の通り。これを1行ずつ実行する。
wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda-repo-wsl-ubuntu-11-7-local_11.7.1-1_amd64.deb
sudo dpkg -i cuda-repo-wsl-ubuntu-11-7-local_11.7.1-1_amd64.deb
sudo cp /var/cuda-repo-wsl-ubuntu-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda完了したら以下コマンドでバージョン確認。
nvcc -Vnvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Jun__8_16:49:14_PDT_2022
Cuda compilation tools, release 11.7, V11.7.99
Build cuda_11.7.r11.7/compiler.31442593_0……となるらしいけど、筆者の場合は Command 'nvcc' not found と言われた。
参考記事から、以下コマンドでパスを通す＆シェルを再起動。
echo 'export PATH="/usr/local/cuda/bin:$PATH"' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH="/usr/local/cuda/lib64:$LD_LIBRARY_PATH"' >> ~/.bashrc
exec "$SHELL"再度バージョン確認すると、今度はうまくいった。
手順5. cuDNNインストール
WSLにCUDA Deep Neural Network (cuDNN) をインストールする。ディープラーニング用にGPUを使うためのNVIDIAのライブラリらしい。
ファイルのダウンロード自体はWindows側でブラウザから行う必要がある。このとき、NVIDIAのデベロッパーアカウントが必要なので作成すること。
以下の公式サイトからアカウント作成とダウンロードを行う。
cuDNN Download | NVIDIA Developer
- 規約を読んで『Agree To the Terms of the cuDNN Software License Agreement』をチェックON。
- CUDAのバージョンを選択。ここではCUDA 11.7なので『Download cuDNN v8.5.0 (August 8th, 2022), for CUDA 11.x』をクリック。
- OSに応じてインストーラをダウンロード。ここではUbuntu 20.04なので『Local Installer for Ubuntu20.04 x86_64 (Deb)』をクリック。
今回ダウンロードしたファイルは cudnn-local-repo-ubuntu2004-8.5.0.96_1.0-1_amd64.deb だった。これをWSL側から実行してインストールする必要がある。
エクスプローラのアドレスバーに \\wsl$ と入力すると、WSLのディストリビューションごとにファイルを閲覧できる。これで当該ファイルをコピペしてホームディレクトリ (/home/{ユーザ名} 以下) に移動させるのが簡単。
その後、以下コマンドで当該ファイルを実行する。
cd ~
sudo dpkg -i cudnn-local-repo-ubuntu2004-8.5.0.96_1.0-1_amd64.deb完了後、最後に以下のメッセージが表示されていた。
The public CUDA GPG key does not appear to be installed.
To install the key, run this command:
sudo cp /var/cudnn-local-repo-ubuntu2004-8.5.0.96/cudnn-local-0579404E-keyring.gpg /usr/share/keyrings/とりあえず言われた通りに以下コマンド実行。これが必要かどうかは不明。
sudo cp /var/cudnn-local-repo-ubuntu2004-8.5.0.96/cudnn-local-0579404E-keyring.gpg /usr/share/keyrings/手順6. Stable Diffusion本体を用意
ようやくStable Diffusionの準備に入る。
Stable Diffusionの本体はGitHubからCloneできる。GitHubのアカウントとアクセストークンは作成済みの前提で進める。
以下コマンドでStable Diffusion本体をホームディレクトリにClone。
cd ~
git clone https://github.com/CompVis/stable-diffusion次に、以下コマンドでAnacondaの仮想環境 ldm を作成する。(この ldm という名前は stable-diffusion/environment.yaml に書いてある。)
cd ~/stable-diffusion
conda env create -f environment.yaml
conda activate ldm……が、筆者の場合は以下のエラーが表示された。
CommandNotFoundError: Your shell has not been properly configured to use 'conda activate'.
To initialize your shell, run
$ conda init <SHELL_NAME>
Currently supported shells are:
- bash
- fish
- tcsh
- xonsh
- zsh
- powershell
See 'conda init --help' for more information and options.
IMPORTANT: You may need to close and restart your shell after running 'conda init'.初めて conda activate をする前に、シェルに応じて conda init を実行してシェルを再起動する必要があるらしい。
Ubuntuならデフォルトのシェルはbashなので、以下コマンドを実行。
conda init bash
exec "$SHELL"するとシェルのユーザ名の前に (base) と表示されるようになった。
もう一度 conda activate ldm を実行すると、シェルのユーザ名の前が (ldm) に変わった。これでOK。
この状態では仮想環境 ldm が有効化されており、stable-diffusion/environment.yaml の dependencies に記載されているライブラリがインストール済みの状態になっている。ここにPyTorch等が含まれている。なお、この状態はコマンド conda deactivate で終了できる。
ここで、念のためPythonからCUDAが使えるか確認する。python コマンドを実行してPythonの対話モードに入り、以下の2行を実行する。
import torch
torch.cuda.is_available()Trueこのように True が表示されれば、PythonからGPUを使用して計算できる環境が整っている。
もし False が表示されてしまった場合、以下を試してみる。
- Windowsを再起動する。
- Windows側でコマンドプロンプトから
wsl --shutdownを実行してWSLを落とし、もう一度起動する。 - ちゃんと環境に合ったNVIDIAドライバーとCUDA Toolkitをインストールしたことを確認する。
ちなみに筆者は、CUDA Toolkitの導入時に間違えてディストリビューション『Ubuntu』を選んでしまって、一回ここで引っかかった。
筆者のようにインストールするものを間違えた場合は、以下を参考にNVIDIA ドライバーとCUDA Toolkitをアンインストールできるかも。
tensorflow - How to remove cuda completely from ubuntu? - Stack Overflow
コメントにはCUDA 11.4以降は /usr/local/cuda-11.4/bin/cuda-uninstaller のようなアンインストーラがあると書いてあるが、11.7では見当たらなかったので、コメント通りに以下コマンドで除去した。
sudo apt-get --purge remove "*cublas*" "cuda*" "nsight*"
sudo apt-get --purge remove "*nvidia*"
sudo rm -rf /usr/local/cuda*手順7. Stable Diffusionの学習モデルを用意
GitHubに置いてあるのはStable Diffusion本体だけで、学習モデルは含まれていない。
学習モデルはHugging Faceに置かれているので、別途ダウンロードする。Hugging Faceのアカウントが必要なので、持っていなければ作成する。(ちなみにアカウント作成の途中で「Captcha failed」で引っかかる人は、IPを変えると通るかも。筆者はスマホからモバイルデータ通信で試したら通った。)
ここではリポジトリのCloneはせずに、ブラウザでファイルを直接ダウンロードする。
以下のページを表示したら、まず上部に表示されている注意事項 (「You need to share your contact information to access this model.」枠内) を読む。
CompVis/stable-diffusion-v-1-4-original · Hugging Face
内容に同意したら、枠内下部のチェックボックスをONにして『Access repository』ボタンを押すと、このリポジトリのコンテンツにアクセスできるようになる。『Access repository』ボタンを押した時点で、リポジトリの管理者にユーザ名とメールアドレスが通知されるので留意。
※内容はちゃんと読むこと。生成した画像の権利はユーザのものだけど、違法だったり有害だったりなものを作らせたり共有したりはダメ。ここで出てくるCreativeML Open RAIL-Mライセンスは学習モデルとその生成物のライセンスらしいので、必要なら調べること。
そしたら、本文の「Download the weights」の下にリンクが書かれている .ckpt ファイルが学習モデル。かなりサイズが大きいので注意。sd-v1-4.ckpt は3.97GB、sd-v1-4-full-ema.ckpt は7.17GBある。違いはよくわからないので、ここでは軽い sd-v1-4.ckpt の方をダウンロードする。
ダウンロードしたファイルを、Stable Diffusion本体のモデル置き場にファイル名 model.ckpt として置く。モデル置き場は以下コマンドで作成する。
mkdir -p ~/stable-diffusion/models/ldm/stable-diffusion-v1これでStable Diffusionのローカルリポジトリ内に作成した models/ldm/stable-diffusion-v1/stable-diffusion-v1 フォルダの中に先程の学習モデルファイルを置き、ファイル名を model.ckpt に変更する。
ホームディレクトリから見て、ファイル stable-diffusion/models/ldm/stable-diffusion-v1/model.ckpt が存在すればOK。
手順8. Stable Diffusion実行
ここまでの手順で準備したStable Diffusionは、以下コマンドで実行できる。
cd ~/stable-diffusion
python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse" --plms……が、筆者環境では途中で以下のエラーが出て終了してしまった。
RuntimeError: CUDA out of memory. Tried to allocate 1.50 GiB (GPU 0; 6.00 GiB total capacity; 4.12 GiB already allocated; 0 bytes free; 4.24 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONFGPUメモリ (VRAM) が足りないらしい。
手順9. Optimized Stable Diffusionを使う場合
冒頭の参考記事に、速度を犠牲にVRAMの使用量を減らしたというフォーク版のOptimized Stable Diffusionがあったので使ってみる。こちらはREADMEにRTX 2060の名前が出ているので動きそう。
必要なのは当該リポジトリ内の optimizedSD 以下のみなので、Cloneして移動のみ行う。
cd ~
git clone https://github.com/basujindal/stable-diffusion.git stable-diffusion-optimized
mv stable-diffusion-optimized/optimizedSD/ stable-diffusion/
rm -rf stable-diffusion-optimized/完了したら、今度は以下コマンドで実行してみる。
cd ~/stable-diffusion
python optimizedSD/optimized_txt2img.py --prompt "radio tower at dusk, look upward" --n_samples 8 --ddim_steps 30 --seed 484601しばらく待つと終了し、生成された画像がStable Diffusion内の outputs/txt2img-samples/ 以下に保存される。
このコマンドは、筆者環境では2～3分程度かかる。使用したオプションは以下の通り。
--prompt… キーワード。ここでは例として、夕暮れの中で見上げた電波塔を指定。--n_samples 8… 画像を8枚出力する。--ddim_steps 30… ステップ数に30回を指定。未指定時は50となる。少なくすると計算回数が減って早くなるが、画像はぐちゃぐちゃになる。--seed 484601… シード値に484601を指定。指定すると、生成される画像をシード値に応じたものに固定できる。
スペックの高くないGPUで遊ぶなら、最初は --ddim_steps を20など少なめに設定して実行してみて、ある程度気に入った画像ができたら --seed で固定しつつ --ddim_steps を上げながら再実行して、結果を見比べるのがいいのかも。
--ddim_steps は単純に増やせば画質が良くなるというわけではなく、変形したり、少ないときはあったものが多いときは消えたりもする。10だと少なすぎるのかノイズが目立つ。
シード値は実行時の出力の1行目に表示されているので、以下のようにそれを見て指定すれば、最初にランダム生成したものでも後から固定可能。
Global seed set to 261508
Loading model from models/ldm/stable-diffusion-v1/model.ckpt
Global Step: 470000
UNet: Running in eps-prediction mode
CondStage: Running in eps-prediction mode
FirstStage: Running in eps-prediction modes.
(以下略)なお、Optimized Stable Diffusionのコマンドラインオプションはオリジナル版と少し違うことを留意。例えば、オリジナル版のREADMEのサンプルに書いてある --plms オプションはこちらには無い。
以下に optimized_txt2img.py のヘルプを以下にコピペしておく。
usage: optimized_txt2img.py [-h] [--prompt [PROMPT]] [--outdir [OUTDIR]] [--skip_grid] [--skip_save] [--ddim_steps DDIM_STEPS] [--fixed_code]
[--ddim_eta DDIM_ETA] [--n_iter N_ITER] [--H H] [--W W] [--C C] [--f F] [--n_samples N_SAMPLES] [--n_rows N_ROWS]
[--scale SCALE] [--device DEVICE] [--from-file FROM_FILE] [--seed SEED] [--unet_bs UNET_BS] [--turbo]
[--precision {full,autocast}] [--format {jpg,png}] [--sampler {ddim,plms}]
optional arguments:
-h, --help show this help message and exit
--prompt [PROMPT] the prompt to render
--outdir [OUTDIR] dir to write results to
--skip_grid do not save a grid, only individual samples. Helpful when evaluating lots of samples
--skip_save do not save individual samples. For speed measurements.
--ddim_steps DDIM_STEPS
number of ddim sampling steps
--fixed_code if enabled, uses the same starting code across samples
--ddim_eta DDIM_ETA ddim eta (eta=0.0 corresponds to deterministic sampling
--n_iter N_ITER sample this often
--H H image height, in pixel space
--W W image width, in pixel space
--C C latent channels
--f F downsampling factor
--n_samples N_SAMPLES
how many samples to produce for each given prompt. A.k.a. batch size
--n_rows N_ROWS rows in the grid (default: n_samples)
--scale SCALE unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))
--device DEVICE specify GPU (cuda/cuda:0/cuda:1/...)
--from-file FROM_FILE
if specified, load prompts from this file
--seed SEED the seed (for reproducible sampling)
--unet_bs UNET_BS Slightly reduces inference time at the expense of high VRAM (value > 1 not recommended )
--turbo Reduces inference time on the expense of 1GB VRAM
--precision {full,autocast}
evaluate at this precision
--format {jpg,png} output image format
--sampler {ddim,plms}
sampler以上。