GPUメモリ8GBでStable Diffusionを動かす

概要

GPUが8GBのマシンでStable Diffusionなモデルを動かした記録．
Windows10, Python3.9.13, CUDA11.2, Torch1.8.1+cu111です．

動かしたもの

waifu-diffusion v1.4 - Diffusion for Weebs　の最小モデル．

結果

・入力　"1girl, simple background, open mouth, cute girl, blond hair, blue eye, upper body."
・出力1 生成時間は20-30秒

・出力2

・入力　"1girl, simple background, close mouth, red hair, black eye, upper body."
・出力1

・出力2

コード

のExample Codeの修正．そのままではエラーとメモリ不足で動かなかったため

• 一部の書き換え
• 最も軽量なモデルの利用．リンク先Open in Colab参照．

import torch
from torch.cuda.amp import autocast
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    'hakurei/waifu-diffusion',
    torch_dtype=torch.float16,
    revision='fp16'
).to('cuda')

prompt = "1girl, simple background, blue hair, blue eye, upper body."
num_samples = 1 # 一度に出力される画像の数
with autocast(True):
    image = pipe([prompt] * num_samples, guidance_scale=7.5)["images"]

# pipe()の出力
# StableDiffusionPipelineOutput(images=[<PIL.Image.Image image mode=RGB size=512x512 at 0x20C34C33FA0>])

image[0].save("test0.png")

以下，細かいこと

エラー1：cannot import name 'autocast' from 'torch'

リンク先のように書き換えればいい．

エラー2：cannot import name 'CLIPImageProcessor' from 'transformers'

既存の環境が原因らしい．仮想環境に入れ直せとのこと．

Pythonの仮想環境

VSCodeがデフォルトでターミナルにPowerSellを使うのでコマンドプロンプトに変更．PowerShellのままだとactivate.batが機能しない．

仮想環境にPytorchインストール

本マシンではRTX3070TiとCUDA 11.2を用いているため，上を参考に

# CUDA 11.1
pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/cu111/torch_stable.html

CUDAのバージョンが異なるがいいらしい．

guidance_scale=について

リンク先が詳しい．見た方が早い．