GTX 1660 SUPERでもStable-Diffusionしたい

# 概要
表題の通りGTX 1660 SUPERのグラボでもStable-Diffusionを動かすための環境構築をやります。
本家のStable-DiffusionはVRAM不足で動かなかったので、有志の方が作ってくださった軽量版のStable-Diffusionを使用します。
https://github.com/basujindal/stable-diffusion

最終的な私の環境での画像生成結果が以下になります。
![seed_888684_00022.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/526299/54af0f12-7993-23de-325e-dc591560c712.png)

![seed_722008_00004.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/526299/4e588a9f-4900-4c65-9245-486b24a39531.png)


軽量版ではありますが、本家と遜色のない結果を得られていることが分かります。ただし、生成にかかる時間は長くなっているようです。
上記の街の画像(512x512)を10枚ほど生成している間のGPUの様子は以下のとおりです。VRAMに1GBほど余裕があることが分かります。生成には10~15分ほどかかります。
```
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.141.03   Driver Version: 470.141.03   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| 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 ...  Off  | 00000000:01:00.0  On |                  N/A |
|100%   81C    P2   121W / 125W |   5060MiB /  5941MiB |     86%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1824      G   /usr/lib/xorg/Xorg                184MiB |
|    0   N/A  N/A     13941      G   /usr/bin/gnome-shell               36MiB |
|    0   N/A  N/A     16615      G   ...198359930652384403,131072       87MiB |
|    0   N/A  N/A     22159      C   python                           4747MiB |


```
こういうAI系を色々と動かす場合は環境を分けたいと思いますので、今回はDockerを使用します。

# 環境
* OS : Ubuntu 22.04.1 LTS
* Docker : version 20.10.18, build b40c2f6
    * Docker Compose : version v2.10.2
* ホストのCUDAバージョン : 11.4
* CPU : Corei5-9400F
* GPU : GTX 1660 SUPER
* RAM : 16GB
* VRAM : 6GB

# 事前準備

DockerでNVIDIAのGPUを使用するには[NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker)をホストの環境にインストールしておく必要があります。

# リポジトリが提供するDockerコンテナを使用する

上記の有志の方が作ってくださったコンテナには `docker-compose.yml` が含まれているので、こちらを使用すればすぐに環境が整ったdockerコンテナを立ち上げることができます。
まずは、リポジトリをクローンしてきます。

```bash
git clone https://github.com/basujindal/stable-diffusion.git
```
続いて、クローンしてきたリポジトリと同じ階層に モデルデータを入れておく`sd-data`フォルダと出力画像が格納される`sd-out`フォルダを作成し、その中にモデルデータをダウンロードします。
モデルデータをダウンロードするためには[Hagging Face](https://huggingface.co/)のアカウントが必要となるので、事前に作成しておきましょう。以下のwgetコマンドの `--http-user`と`--http-passwd`にはHagging Faceに登録したアカウントのメールアドレスとパスワードを設定します。
```bash
mkdir sd-data
mkdir sd-out
cd sd-data
wget --http-user=hoge --http-passwd=password https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt
```
ダウンロードが完了したら、モデルファイルのファイル名を `model.ckpt`に変更します。
```bash
mv sd-v1-4.ckpt model.ckpt
```
この時点でのフォルダ構成は以下のようになっているはずです。
```
.
├── sd-data
│   └── model.ckpt
├── sd-output
└── stable-diffusion
    ├── assets
    ...
```

続いて、Docker Composeをインストールします。インストール済みの場合はスキップしてください。
```bash
sudo apt update
sudo apt install docker-compose-plugin
```
正常にインストールができていれば、以下のコマンドでバージョンを確認できます。
```bash
docker compose version
(output) Docker Compose version v2.10.2
```
composeがインストールできたらクローンしてきたリポジトリの中に入ってコンテナを立ち上げます。
```bash
cd stable-diffusion
docker compose up --build
```
2回目以降の起動時は `--build` オプションは不要です。
```bash
docker compose up
```
このイメージは自動的にGradioというウェブアプリを立ち上げるようになっており、http://localhost:7860/ にアクセスするとGUIで画像生成が出来るようになっています。
![Screenshot from 2022-09-24 16-19-49.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/526299/23d305f3-f17f-aa4e-8524-b494d6e647e4.png)

なお、GTX世代のGPUを使用している人は、一番下にある `full_precision`のチェックボックスを**必ずオンにしてください**。
デフォルトでは半精度小数点(half)を利用しているらしいのですが、GTX世代では対応していないので、 `full_precision` (float)にする必要があるそうです。

出力した画像は自動的に `sd-output` フォルダの中に保存されます。

## CUIから操作する
GUIではなく、直接スクリプトを叩きたい場合もあるかと思います。
そのような場合、デフォルトの状態でコンテナの中に入ってスクリプトを実行すると恐らくPCがフリーズするかコンテナが落ちてしまいます。理由はコンテナ起動時に自動的に立ち上がるGradioがすでにモデルをロードしておりVRAMを消費しているためで、その状態で更にPythonスクリプトを叩いてモデルをロードするとVRAM不足になってしまいます。

よって、まずは起動時にGradioのプロセスが立ち上がらないようにします。
クローンしてきたリポジトリの中の`Dockerfile`の最後の行をコメントアウトします。
```Dockerfile:Dockerfile
FROM nvidia/cuda:11.3.1-base-ubuntu20.04

SHELL ["/bin/bash", "-c"]

RUN apt update \
 && apt install --no-install-recommends -y curl wget git \
 && apt-get clean

RUN wget https://repo.anaconda.com/miniconda/Miniconda3-py38_4.12.0-Linux-x86_64.sh -O ~/miniconda.sh \
 && bash ~/miniconda.sh -b -p $HOME/miniconda \
 && $HOME/miniconda/bin/conda init

COPY . /root/stable-diffusion

RUN eval "$($HOME/miniconda/bin/conda shell.bash hook)" \
 && cd /root/stable-diffusion \
 && conda env create -f environment.yaml \
 && conda activate ldm \
 && pip install gradio==3.1.7

VOLUME /root/.cache
VOLUME /data
VOLUME /output

ENV PYTHONUNBUFFERED=1
ENV GRADIO_SERVER_NAME=0.0.0.0
ENV GRADIO_SERVER_PORT=7860
EXPOSE 7860

RUN ln -s /data /root/stable-diffusion/models/ldm/stable-diffusion-v1 \
 && mkdir -p /output /root/stable-diffusion/outputs \
 && ln -s /output /root/stable-diffusion/outputs/txt2img-samples

WORKDIR /root/stable-diffusion

ENTRYPOINT ["/root/stable-diffusion/docker-bootstrap.sh"]
# CMD python optimizedSD/txt2img_gradio.py この行をコメントアウト
```
その後、 `docker-compose.yml` も以下のように編集します。
```yaml:docker-compose.yml
version: "3.9"
services:
  sd:
    build: .
    ports:
      - "7860:7860"
    volumes:
      - ../sd-data:/data
      - ../sd-output:/output
      - sd-cache:/root/.cache
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    command: tail -f /dev/null ← この行を追加
volumes:
  sd-cache:
```
この行を追加しないとコンテナが立ち上がった瞬間に終了してしまいます。いつまでも終わらない処理であれば何でもいいと思います。コンテナの永続化について検索すると `tty: true` を追加するという方法もヒットしますが、私の環境では効果がありませんでした。

以上の編集が出来たら、再度 `docker compose up --build` を実行しイメージをビルドします。2回目以降の実行では `--build`オプションは必要ありません。

コンテナが立ち上がったら中に入ります。

```bash
sudo docker exec -it stable-diffusion-optimized-sd-1 /bin/bash
```
まずはコンテナに事前に用意されたconda環境を立ち上げます。
```bash
conda activate ldm
```
この環境でないと `torch` 等の必要なパッケージが存在せずPythonスクリプトを実行できないので注意してください。
テキストから画像を生成する際は以下のコマンドで生成することができます。各引数の意味はリポジトリのREADMEを参照してください。
GTX世代のGPUを使用する場合は、Gradioで`precision_full`にチェックを入れたように、一番最後の `--precision full`の引数を忘れないでください。
```bash
python optimizedSD/optimized_txt2img.py --prompt "Robot Ninja" --H 512 --W 512 --n_iter 2 --n_samples 10 --ddim_steps 50 --precision full
```
![seed_793827_00004.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/526299/cfed88c9-d9a3-f833-e068-327ab4b867eb.png)

Gradioを立ち上げたい場合は、`txt2img_gradio.py`を実行してください。先ほどと同様にhttp://localhost:7860/ にアクセスすればGradioを使用することが出来ます。
```bash
python optimizedSD/txt2img_gradio.py
```

# 一から環境を作る場合

上記の手順は全て有志の方のリポジトリに用意されているDockerイメージを使用しましたが、そちらを使わずに自分で一つずつ必要なライブラリ等をインストールして環境構築を行う手順を紹介します。
...が、基本的には上記のDockerイメージをそのまま使えば問題ないはずですし、何か改造をしたい場合や、本家のStable-Diffusionのイメージを作りたい場合でも、やはり上記のDockerイメージをベースに作業をした方が早いと思います。そもそもDockerfileを見れば何をすればいいか全部書いてあるし...

実は自分は軽量版のリポジトリにDockerfileが用意されていることに気づく前に自分で環境構築をやってしまいました。そちらでも問題なく動かすことが出来たので、せっかくですし何かの参考になるかもしれないので一応書いておきます。

まずはベースとなるコンテナを立ち上げます。イメージには以下のnvidiaのcudaコンテナを使用します。ここで、引数に `--gpus all`を付けるのを忘れないでください。これが無いとコンテナの中でGPUが使えません。
```bash
sudo docker run --gpus all -it -d --name stable-diffusion nvidia/cuda:11.3.1-base-ubuntu20.04
```
ここで`docker: Error response from daemon: pull access denied for repository does not exist or may require ‘docker login’`というエラーが出た場合は `sudo docker login`を実行してdockerhubへのログイン情報を入力してやる必要があります。この時に求められるユーザ名はdockerhubに登録した**ユーザ名**である必要があり、**メールアドレス**では駄目なので注意が必要です。
コンテナが立ち上がったら中に入ります。
```bash
sudo docker container exec -it stable-diffusion /bin/bash
```

## Pythonのインストール
まずはPythonをインストールします。`apt`からインストールすると`pip`を使えなかったので、ソースコードからインストールします。まずはそのために必要なツールをインストールします。
```bash
apt update
apt install build-essential zlib1g-dev libncurses5-dev libgdbm-dev libnss3-dev libssl-dev libreadline-dev libffi-dev libbz2-dev liblzma-dev wget software-properties-common git
```
続いてPythonのソースコードをダウンロードし、解凍します。ここでは3.8.10を使用します。
```bash
wget https://www.python.org/ftp/python/3.8.10/Python-3.8.10.tgz
tar -xf Python-3.8.10.tgz
```
解凍したフォルダの中に入り、`configure`を実行します。
```bash
cd Python-3.8.10
./configure --enable-optimizations
```
configureが完了したら`make`を実行し、ビルドを行います。
```bash
make -j $(nproc)
```
ビルドが完了したらインストールします。
```bash
make altinstall
```

## Pythonライブラリのインストール
Pythonがインストールできたら、Stable-Diffusionが依存しているPythonライブラリをインストールします。まずは`torch`をインストールします。
```bash
python3.8 -m pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
```
続いて、その他の必要なライブラリをインストールします。
有志の方が[GitHub](https://raw.githubusercontent.com/dai-ichiro/env4stable-diffusion/main/requirements.txt)で`requirements.txt`としてまとめて下さっているので、そちらを使用します。
```bash
python3.8 -m pip install -r https://raw.githubusercontent.com/dai-ichiro/env4stable-diffusion/main/requirements.txt
```

## Optimized Stable-Diffusionのダウンロード
ここまでで環境構築は完了しているので、肝心の軽量化版Stable-Diffusionをダウンロードします。
```
git clone https://github.com/basujindal/stable-diffusion.git
```
続いて、モデルファイルをダウンロードします。その前にモデルを配置するためのフォルダを用意します。
```bash
cd models/ldm
mkdir stable-diffusion-v1
cd stable-diffusion-v1
```
クローンしてきた`stable-diffusion`フォルダの下にHugging Faceからダウンロードしてきたモデルファイルを配置します。
上にも書いてあるように、`--http-user`と`--http-password`はhaggingfaceのアカウントのユーザ名とパスワードにしてください。ここではユーザ名としてメールアドレスを使用することが出来ます。
```bash
wget --http-user=hoge --http-passwd=password https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt
```
ダウンロードしたモデルのファイル名を`model.ckpt`に変更します。
```bashrc
mv sd-v1-4.ckpt stable-diffusion/sd-v1-4.ckpt
```
続いて、クローンしてきたリポジトリの中のモジュールを読み込めるように環境変数を設定します。
```bash
export PYTHONPATH="$PYTHONPATH:/stable-diffusion"
```
また、この設定を永続化するために`.bashrc`に追加します。
```bash
echo 'export PYTHONPATH="$PYTHONPATH:/home/test/lib"' >> ~/.bashrc
```
ここまで設定できれば動作するようになっているはずです！
```bash
python3.8 optimizedSD/optimized_txt2img.py --prompt "mountain turtle" --H 512 --W 512 --n_iter 2 --n_samples 1 --ddim_steps 50 --precision full
```
出力はstable-diffusionフォルダ以下の`outputs/txt2img-samples/mountain_turtle`フォルダに保存されているはずです。
画像をホストOSに持ってくるためにはホスト側でdockerのcpコマンドを使用します。
```bash
sudo docker cp stable-diffusion:/stable-diffusion/outputs/txt2img-samples/mountain_turtle/<filename> <filename>
```
`<filename>`となっている部分には画像ファイルのファイル名を入れてください。
![seed_748248_00000.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/526299/95ff8dba-52a9-dca2-20fb-1acf176b659f.png)


## Dockerコンテナのイメージ化
最後にここまで作成したコンテナを再利用できるようにイメージとしてコミットします。
```bash
sudo docker commit stable-diffusion <作成者の名前>/<イメージ名>:<タグ>
```

と、ここまで色々とやってきましたが、最初にも述べたように軽量版のリポジトリにすでに用意されているDockerfileをそのまま利用するほうが確実だし手っ取り早いと思います。Dockerではなく仮想マシンやホストOSに環境構築をしたい場合等はここの情報が役に立つかもしれませんが、その際はここで紹介した手順はnvidiaのcudaコンテナのイメージをベースにしているため、普通のUbuntu20.04とは異なる可能性がある点に注意してください。

なお、上記の環境で本家のStable-Diffusionも動かすことが出来ると思いますが、私の環境では途中で`CUDA run out of memory`してしまうため最終的な動作確認は取れていません。