GPU メモリ 8GB で StableDiffusion を動かす

概要

StableDiffusion のデフォルト設定がメモリが少ないマシンだと動かなかったためやったことをまとめました。

環境

FROM nvcr.io/nvidia/pytorch:22.04-py3

RUN pip install diffusers==0.2.4 transformers==4.21.2 ftfy==6.1.1

やったこと

モデルを半精度化する

まず8GBだと単精度だとモデルが GPU に乗らないので読み込み時に半精度で読み込みます。単精度時の　StableDiffusion の重みが10GBくらいだったのでこれで5GBになり、メモリに乗ります。なお、CPU で動かそうとすると無限時間かかるので基本的に GPU に乗せることは must になります。

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    'CompVis/stable-diffusion-v1-4',
    use_auth_token=TOKEN, 
    revision='fp16',
    torch_dtype=torch.float16,
)
pipe.to('cuda')

TOKEN は huggingface で発行したトークンを入れてください。

mixed precision を使って入力を半精度に cast する

StableDiffusionPipeline はデフォルトだと入力を半精度に cast してくれない¹ので推論時に mixed precision のコンテキストブロックを作って半精度に cast します。これで推論時エラーは起きません。

with torch.autocast('cuda'):
    image = pipe(
        PROMPT,
        height=512,
        width=512,
        guidance_scale=7.5,
        num_inference_steps=50,
    )['sample'][0]

PROMPT は生成呪文を入れてください。

バッチサイズは1にする

試しにバッチサイズを上げてみたらメモリエラーで死んだのでバッチサイズは1のみです。メモリが12GBあったらバッチサイズもう少し上げれるかもしれません。バッチサイズを上げる場合は [PROMPT] * batch_size のようにリスト化＋バッチサイズの長さにコピーすれば並列化の恩恵にあずかれます。

その他メモ

• 半精度化 (mixed precision) しても生成クオリティは落ちない。推論速度が上がるためやりトク
• 生成サイズはデフォルトから落とすと生成クオリティが落ちるらしい (友人談)
• num_inference_steps はデフォルトから上げても正直そこまで変わらないので固定でよさそう
  • 100にしても50と比べて定性的には特に良くならない
• StableDiffusion のうまくいっているサンプルを見た感じ、生成したい画像の概要＋生成クオリティを上げるコマンド連打がいい感じのクオリティの画像を作るコツっぽい

1. そのまま inference すると型エラーが出ます。 ↩