概要

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

最終的な私の環境での画像生成結果が以下になります。

軽量版ではありますが、本家と遜色のない結果を得られていることが分かります。ただし、生成にかかる時間は長くなっているようです。
上記の街の画像(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をホストの環境にインストールしておく必要があります。

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

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

git clone https://github.com/basujindal/stable-diffusion.git

続いて、クローンしてきたリポジトリと同じ階層にモデルデータを入れておくsd-dataフォルダと出力画像が格納されるsd-outフォルダを作成し、その中にモデルデータをダウンロードします。
モデルデータをダウンロードするためにはHagging Faceのアカウントが必要となるので、事前に作成しておきましょう。以下のwgetコマンドの --http-userと--http-passwdにはHagging Faceに登録したアカウントのメールアドレスとパスワードを設定します。

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に変更します。

mv sd-v1-4.ckpt model.ckpt

この時点でのフォルダ構成は以下のようになっているはずです。

.
├── sd-data
│   └── model.ckpt
├── sd-output
└── stable-diffusion
    ├── assets
    ...

続いて、Docker Composeをインストールします。インストール済みの場合はスキップしてください。

sudo apt update
sudo apt install docker-compose-plugin

正常にインストールができていれば、以下のコマンドでバージョンを確認できます。

docker compose version
(output) Docker Compose version v2.10.2

composeがインストールできたらクローンしてきたリポジトリの中に入ってコンテナを立ち上げます。

cd stable-diffusion
docker compose up --build

2回目以降の起動時は --build オプションは不要です。

docker compose up

このイメージは自動的にGradioというウェブアプリを立ち上げるようになっており、http://localhost:7860/ にアクセスするとGUIで画像生成が出来るようになっています。

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

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

CUIから操作する

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

よって、まずは起動時にGradioのプロセスが立ち上がらないようにします。
クローンしてきたリポジトリの中の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 も以下のように編集します。

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オプションは必要ありません。

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

sudo docker exec -it stable-diffusion-optimized-sd-1 /bin/bash

まずはコンテナに事前に用意されたconda環境を立ち上げます。

conda activate ldm

この環境でないと torch 等の必要なパッケージが存在せずPythonスクリプトを実行できないので注意してください。
テキストから画像を生成する際は以下のコマンドで生成することができます。各引数の意味はリポジトリのREADMEを参照してください。
GTX世代のGPUを使用する場合は、Gradioでprecision_fullにチェックを入れたように、一番最後の --precision fullの引数を忘れないでください。

python optimizedSD/optimized_txt2img.py --prompt "Robot Ninja" --H 512 --W 512 --n_iter 2 --n_samples 10 --ddim_steps 50 --precision full

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

python optimizedSD/txt2img_gradio.py

一から環境を作る場合

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

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

まずはベースとなるコンテナを立ち上げます。イメージには以下のnvidiaのcudaコンテナを使用します。ここで、引数に --gpus allを付けるのを忘れないでください。これが無いとコンテナの中でGPUが使えません。

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に登録したユーザ名である必要があり、メールアドレスでは駄目なので注意が必要です。
コンテナが立ち上がったら中に入ります。

sudo docker container exec -it stable-diffusion /bin/bash

Pythonのインストール

まずはPythonをインストールします。aptからインストールするとpipを使えなかったので、ソースコードからインストールします。まずはそのために必要なツールをインストールします。

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を使用します。

wget https://www.python.org/ftp/python/3.8.10/Python-3.8.10.tgz
tar -xf Python-3.8.10.tgz

解凍したフォルダの中に入り、configureを実行します。

cd Python-3.8.10
./configure --enable-optimizations

configureが完了したらmakeを実行し、ビルドを行います。

make -j $(nproc)

ビルドが完了したらインストールします。

make altinstall

Pythonライブラリのインストール

Pythonがインストールできたら、Stable-Diffusionが依存しているPythonライブラリをインストールします。まずはtorchをインストールします。

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でrequirements.txtとしてまとめて下さっているので、そちらを使用します。

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

続いて、モデルファイルをダウンロードします。その前にモデルを配置するためのフォルダを用意します。

cd models/ldm
mkdir stable-diffusion-v1
cd stable-diffusion-v1

クローンしてきたstable-diffusionフォルダの下にHugging Faceからダウンロードしてきたモデルファイルを配置します。
上にも書いてあるように、--http-userと--http-passwordはhaggingfaceのアカウントのユーザ名とパスワードにしてください。ここではユーザ名としてメールアドレスを使用することが出来ます。

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に変更します。

mv sd-v1-4.ckpt stable-diffusion/sd-v1-4.ckpt

続いて、クローンしてきたリポジトリの中のモジュールを読み込めるように環境変数を設定します。

export PYTHONPATH="$PYTHONPATH:/stable-diffusion"

また、この設定を永続化するために.bashrcに追加します。

echo 'export PYTHONPATH="$PYTHONPATH:/home/test/lib"' >> ~/.bashrc

ここまで設定できれば動作するようになっているはずです！

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コマンドを使用します。

sudo docker cp stable-diffusion:/stable-diffusion/outputs/txt2img-samples/mountain_turtle/<filename> <filename>

<filename>となっている部分には画像ファイルのファイル名を入れてください。

Dockerコンテナのイメージ化

最後にここまで作成したコンテナを再利用できるようにイメージとしてコミットします。

sudo docker commit stable-diffusion <作成者の名前>/<イメージ名>:<タグ>

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

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

GTX 1660 SUPERでもStable-Diffusionしたい

概要

環境