DockerでStable diffusionを動かす

WindowsでDockerからStable Diffusionを実行したいけどメモリが足りないから色々調べた。

対象者

  • Windows環境
  • Stable Diffusionの概要を知っている
  • Docker経験者
  • ホストOS(Windows)を汚したくない
  • グラフィックボードを搭載している
  • 強力なグラフィックボードを持っていない

なぜDockerで動かす必要がある?

  • Web上でStable Diffusionを試すことができるがなんだかんだで遅い。
  • Pythonを使ってないため、ホストOSで環境構築を行うのが嫌。

それで、Dockerなら使い捨てでコンテナとイメージを削除すればきれいさっぱりすることができると思った。

検証環境

恐らくCPUやメモリのスペックはそこまで必要でない印象だが、グラフィックボードの搭載メモリ10GB以上が推奨となっている。

私の所持しているグラボはRTX 2080Superのメモリ8GBのため、実行時にメモリが足らずにエラーが発生するがtorch.float16にすることによって、メモリ不足エラーを回避する。対処法は後に記述。

事前準備

https://huggingface.co/にアクセスし、アカウントを作成しておく。メールアドレスと名前の設定のみで、クレジットカードなどの決済情報の入力はない。メールが来るのでリンクをクリックして終了。

https://huggingface.co/CompVis/stable-diffusion-v-1-4-originalにアクセスし、Download the weightsのsd-v1-4.ckptファイル(約4GB)をダウンロードしておく。

Dockerfile

Dockerfileは下記の通り、同じディレクトリに、ダウンロードしたsd-v1-4.ckptを置いておくように。

FROM continuumio/miniconda3
RUN apt-get update -y && apt-get upgrade -y \
    && apt-get install -y git curl gnupg vim libpng-dev
RUN distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
    && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
    && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
RUN apt-get update -y && apt-get install -y nvidia-docker2
RUN git clone https://github.com/CompVis/stable-diffusion.git /app
RUN cd /app \
    && conda env create -f environment.yaml 
RUN mkdir -p /app/models/ldm/stable-diffusion-v1
COPY sd-v1-4.ckpt /app/models/ldm/stable-diffusion-v1/model.ckpt
CMD [ "sh" ]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

ビルドしておく、タグ名は適当

docker build . -t hoge:latest
1

dockerからホストOSのCPUを利用するには--gpus allオプションが必要となる。こんな感じ。

docker run -it --gpus all hoge:latest bash
1

torch.float16のスクリプトの用意

ここはコンテナ内の作業

グラフィックボードが推奨に満たしていないため、torch.float16のスクリプトを作成する。

  • txt2img.pyをコピーし、txt2imgf16.pyを作成する。
  • txt2imgf16.pyにtorch.float16設定を記述する。
$ cp /app/scripts/txt2img.py /app/scripts/txt2imgf16.py 
1

編集

$ vim /app/scripts/txt2imgf16.py 
1

model = model.to(torch.float16)を追加する。(240行あたり)

    config = OmegaConf.load(f"{opt.config}")
    model = load_model_from_config(config, f"{opt.ckpt}")
    model = model.to(torch.float16)
1
2
3

Stable Diffusionを実行する

$ conda activate ldm
$ cd /app
$ python scripts/txt2imgf16.py --prompt "Frustration of men seeking marriage" --plms  --H 512 --W 512 --n_samples 1
1
2
3

実行中はダイアログがずらずら流れて、成功すると最後にEnjoyが出てくる。

画像は/app/outputs/txt2img-samplesディレクトリに出力されるので、VS Codeでattachするなり、docker runにvolumeを指定するなりすればよい。

参考リンク