WSL2で今話題のStable Diffusionを動かしてみた

まずはUbuntuを最新の状態にアップデート
```
＄ sudo apt update
$ sudo apt upgrade
```

NVIDIA　GPU(CUDA）が使える状態になっている事を確認
```console:
$ nvidia-smi
Mon Oct  3 11:03:01 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:09:00.0  On |                  N/A |
|  0%   52C    P0    47W / 170W |   1478MiB / 12288MiB |      3%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A        24      G   /Xwayland                       N/A      |
+-----------------------------------------------------------------------------+
```

[NVIDIA Developerのサイト](https://developer.nvidia.com/)からCUDA　Toolkitをダウンロード
(ダウンロードにはNVIDIA Developerへの登録が必要)
[cuDNN 8.5.0 Installation Guide](https://docs.nvidia.com/deeplearning/cudnn/install-guide/index.html)

ダウンロードしたローカルリポジトリのパッケージをインストールし、追加されたNVIDIAのリポジトリからnvidia-cuda-toolkitをインストール
```bash:
$ sudo dpkg -i cudnn-local-repo-ubuntu2004-8.5.0.96_1.0-1_amd64.deb
$ sudo cp /var/cudnn-local-repo-ubuntu2004-8.5.0.96/cudnn-local-0579404E-keyring.gpg /usr/share/keyrings/
$ sudo apt update
$ sudo apt install nvidia-cuda-toolkit
```

インストールされたCUDA　Toolkitのバージョンを確認
```bash:
$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Sun_Jul_28_19:07:16_PDT_2019
Cuda compilation tools, release 10.1, V10.1.243
```

Anacondaのインストールスクリプトをダウンロード
```bash:
$ wget https://repo.anaconda.com/archive/Anaconda3-2022.05-Linux-x86_64.sh
```

シェルスクリプトでAnacondaをインストール
シェルスクリプト`Anaconda3-2022.05-Linux-x86_64.sh`をインストールするといろいろ聞かれるけど、ほぼデフォルトの選択肢でOK
```bash:
$ sh ./Anaconda3-2022.05-Linux-x86_64.sh
```

まずはライセンスの内容が表示されるので「Enter」キーで内容を確認
```console:
Welcome to Anaconda3 2022.05

In order to continue the installation process, please review the license
agreement.
Please, press ENTER to continue
>>>
```

ライセンスの内容をよく読んで規約に同意できるなら「yes」を入力して「Enter」
```console:
The following packages listed on https://www.anaconda.com/cryptography are inclu
ded in the repository accessible through Anaconda Distribution that relate to cr
yptography.

Last updated February 25, 2022

Do you accept the license terms? [yes|no]
[no] >>> yes
```

Anacondaをどこにインストールするか(通常はユーザーのホームディレクトリのanaconda3)にインストールするか尋ねられるので、変更が必要なければそのまま「Enter」
```console:
Anaconda3 will now be installed into this location:
/home/user/anaconda3

  - Press ENTER to confirm the location
  - Press CTRL-C to abort the installation
  - Or specify a different location below

[/home/user/anaconda3] >>>
```

Anaconda3のインストールが完了するとIntelのScikit-learn Extension（機械学習用アクセラレータ)をインストールするかを尋ねられるけど、今回はNVIDIAのCUDAを使うのでこれは無視。
最後にAnaconda３を今ここで初期化するか聞かれるのでこれはご自由に。
ここで「no」を答えても後から手動で初期化の作業は行えますが「yes」を答えると今利用しているコマンドシェルの環境に合わせて環境初期化のための追加スクリプトが追加され、Anaconda3の環境をセットアップしてくれます。
```console:
Preparing transaction: done
Executing transaction: \

    Installed package of scikit-learn can be accelerated using scikit-learn-intelex.
    More details are available here: https://intel.github.io/scikit-learn-intelex

    For example:

        $ conda install scikit-learn-intelex
        $ python -m sklearnex my_application.py



done
installation finished.
Do you wish the installer to initialize Anaconda3
by running conda init? [yes|no]
[no] >>> no
```

ここで「no」と答えた場合には以下のように後ほど手動でAnaconda3の初期化を設定する方法が示されます。
```console:
You have chosen to not have conda modify your shell scripts at all.
To activate conda's base environment in your current shell session:

eval "$(/home/user/anaconda3/bin/conda shell.YOUR_SHELL_NAME hook)"

To install conda's shell functions for easier access, first activate, then:

conda init

If you'd prefer that conda's base environment not be activated on startup,
   set the auto_activate_base parameter to false:

conda config --set auto_activate_base false

Thank you for installing Anaconda3!
```

例えばUbuntuデフォルトのBashを使っているなら以下のように実行すればAnaconda3を利用するのに必要な設定が適用されます。
```bash:
$ eval "$(/home/user/anaconda3/bin/conda shell.bash hook)"
```

```bash:
$ conda init
no change     /home/user/anaconda3/condabin/conda
no change     /home/user/anaconda3/bin/conda
no change     /home/user/anaconda3/bin/conda-env
no change     /home/user/anaconda3/bin/activate
no change     /home/user/anaconda3/bin/deactivate
no change     /home/user/anaconda3/etc/profile.d/conda.sh
no change     /home/user/anaconda3/etc/fish/conf.d/conda.fish
no change     /home/user/anaconda3/shell/condabin/Conda.psm1
no change     /home/user/anaconda3/shell/condabin/conda-hook.ps1
no change     /home/user/anaconda3/lib/python3.9/site-packages/xontrib/conda.xsh
no change     /home/user/anaconda3/etc/profile.d/conda.csh
modified      /home/user/.bashrc

==> For changes to take effect, close and re-open your current shell. <==
```

デフォルトではユーザーのログイン時にAnacondaのbase環境が自動的に選択されるように設定されますが、これが不要な場合には以下のようにこれを禁止することもできます。
(この場合にはAnaconda3の隔離された環境を利用したい場合には毎回conda actibateで必要な環境をアクティベートする必要があります)
```
$ conda config --set auto_activate_base false
```

Anaconda３の環境を最新の状態にアップデートしておく
```
$ conda update -n base -c defaults conda
```

Stable Diffusionのソースコードを取得し`conda`でldm(Latent Diffusion Models)環境を構築
```console:
git clone https://github.com/CompVis/stable-diffusion.git
cd stable-diffusion/
conda env create -f environment.yaml
```

pipの依存関係のインストールにすごく時間がかかる！
(多分こんな表示で延々待たされるがじっと我慢)
```console:
done
Installing pip dependencies: |
```

最後にこんな表示が出たら環境の構築は終了
```console:
done
#
# To activate this environment, use
#
#     $ conda activate ldm
#
# To deactivate an active environment, use
#
#     $ conda deactivate
```

でも、この状態だとWSL2のCUDAドライバーと互換性がなかったり足りないモジュールがあるらしいのでAnaconda３の環境をldmに切り替えてからこれらを追加でインストールする
```bash:
$ conda activate ldm
$ pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
$ pip install transformers==4.19.2 diffusers invisible-watermark
$ pip install -e .
```

以下のようなtorchからcudaの利用可否を確認するPythonスクリプトを作ってこれを実行してみる
```python:check_cuda.py
import torch

print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.device_count())
print(torch.cuda.current_device())
print(torch.cuda.get_device_name())
print(torch.cuda.get_device_name(torch.device('cuda:0')))
print(torch.cuda.get_device_name('cuda:0'))
```

以下のようにCUDAのバージョンやGPUの名前が表示されれば成功
```bash:
$ python check_cuda.py
1.12.1+cu116
True
1
0
NVIDIA GeForce RTX 3090
NVIDIA GeForce RTX 3090
NVIDIA GeForce RTX 3090
```

aptでリポジトリからgit-lfs(gitで大きなファイルを使うためのプラグイン）をインストール
```bash:
$ sudo apt install git-lfs
$ git lfs install
```

HuggingFaceのリポジトリから学習済みのチェックポイントをクローンする
(HuggingFaceからのリポジトリのクローンにはHuggingFaceのアカウントとリポジトリへのアクセス権が必要なので事前に取得しておく)
```
$ git clone https://huggingface.co/CompVis/stable-diffusion-v-1-4-original
Username for 'https://huggingface.co': user@mail.com
Password for 'https://user@mail.com@huggingface.co':
...
```

2回目のパスワード入力後はかなり大きなファイルを取得するらしくかなりの時間待たされるので我慢
```console:
Username for 'https://huggingface.co': user@mail.com
Password for 'https://user@mail.com@huggingface.co':
...
```

チェックポイントの取得が終わったらStable DiffusionのドキュメントにあるようにチェックポイントへのシンボリックリンクをStable Diffusionのリポジトリに作成しておく。
（この時に作成するインボリックリンクは絶対パスで作らないとスクリプトの実行に失敗するので注意)
```bash:
$ cd stable-diffusion
$ mkdir -p models/ldm/stable-diffusion-v1
$ ln -sf /Users/Shared/Develop/Python/stable-diffusion-v-1-4-original/sd-v1-4.ckpt models/ldm/stable-diffusion-v1/model.ckpt
```

サンプルのスクリプトを実行してみる
```
$ python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse" --plms
```

何故かcondaで指定されている依存関係のいくつかがインストールされていなかったり互換性がなかったりするので追加インストール
```bash:
$ pip install omegaconf
$ pip install einops
$ pip install pytorch_lightning
$ pip install taming-transformers
$ pip install taming-transformers-rom1504
$ pip install clip
$ pip install kornia
```

もう一度サンプルのスクリプトを実行してみる
```
$ python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse" --plms
```

結果として`outputs/txt2img-samples`いかに画像ファイルが出力されていれば無事環境構築は完了！
```console:
$ find outputs/txt2img-samples
outputs/txt2img-samples
outputs/txt2img-samples/grid-0000.png
outputs/txt2img-samples/samples
outputs/txt2img-samples/samples/00000.png
outputs/txt2img-samples/samples/00001.png
outputs/txt2img-samples/samples/00002.png
outputs/txt2img-samples/samples/00003.png
outputs/txt2img-samples/samples/00004.png
outputs/txt2img-samples/samples/00005.png
```