【Diffusers】VRAM 6GBしかないノートPCで「Stable Diffusion 2.1-base」をローカルで動かす

公開日:2022年12月17日
最終更新日:2023年2月11日

はじめに

Stable Diffusionをローカルで実行するためにはGitHubリポジトリをダウンロードする方法とHugging Faceリポジトリをダウンロードする方法の2つがあるようです。

Hugging Faceリポジトリをダウンロードする方法

今回はこちらの方法について書きます。

「Diffusers」というライブラリを使います。

GitHubリポジトリをダウンロードする方法

GitHubリポジトリをダウンロードしてStable Diffusionを使う方法はこちらです。
touch-sp.hatenablog.com
こちらの方法はVRAM 6GBでは実行できませんでした。

PC環境

Windows 11
RTX 3060 Laptop (VRAM 6GB)
CUDA 11.6.2
Git for Windows 2.39.0
Python 3.10.9

環境構築

pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
pip install diffusers==0.12.1
pip install transformers==4.26.1 accelerate==0.16.0 scipy==1.10.0
pip install xformers==0.0.17.dev447
pip install safetensors==0.2.8

他のモデルを使用する場合には新しいdiffusersやtransformersを要求される場合があります。

pip install git+https://github.com/huggingface/diffusers.git
pip install git+https://github.com/huggingface/transformers.git

Hugging Faceのリポジトリダウンロード

初回のみこの作業が必要になります。

方法1

git lfs install
git clone https://huggingface.co/stabilityai/stable-diffusion-2-1-base

これで学習済みパラメーターも一緒にダウンロードされます。

方法2

from huggingface_hub import snapshot_download
snapshot_download(repo_id="stabilityai/stable-diffusion-2-1-base", revision="main")



方法2の方が速いような気がします。

Pythonスクリプト

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

model_id = "./stable-diffusion-2-1-base"

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")
pipe.enable_attention_slicing()

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]
    
image.save("astronaut_rides_horse.png")



以下のようなエラーがでますが動作には問題なさそうです。

A matching Triton is not available, some optimizations will not be enabled.
Error caught was: No module named 'triton'



2023年1月現在Windowsでは「triton」は使えません。
WSL2で「triton」ある、なしを比較しましたがVRAM使用量には大差ない印象です。

Windowsユーザーは無視して問題ないと考えます。

少ないVRAMで動かすために必要なこと

  • xFormersを使う(Windowsではtriton使えないけど問題なし)
  • torch.float16を使う
  • enable_attention_slicing()を使う

拡張

seedを使いたい

上記スクリプトを実行すると同じpromptでも毎回結果が異なります。

結果を再現したい場合にはseedを設定する必要があります。

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

model_id = "./stable-diffusion-2-1-base"

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")
pipe.enable_attention_slicing()

prompt = "a photo of an astronaut riding a horse on mars"
seed = 200

generator = torch.manual_seed(seed)
image = pipe(prompt, generator=generator).images[0]
    
image.save("astronaut_rides_horse.png")

出力を複数にしたい

上記スクリプトでは出力は1画像のみです。

複数にしたい場合には num_images_per_prompt を設定します。

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

model_id = "./stable-diffusion-2-1-base"

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")
pipe.enable_attention_slicing()

prompt = "a photo of an astronaut riding a horse on mars"
seed = 200
num_images_per_prompt = 5

generator = torch.manual_seed(seed)
images = pipe(
    prompt = prompt,
    generator = generator,
    num_images_per_prompt = num_images_per_prompt).images

for i, image in enumerate(images):
    image.save(f"result_{i}.png")

ネガティブプロンプトを使いたい

ネガティブプロンプトを使うのも非常に簡単です。

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

model_id = "./stable-diffusion-2-1-base"

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")
pipe.enable_attention_slicing()

prompt = "anime of tsundere moe kawaii beautiful girl"
negative_prompt = "red eyes red hair"
seed = 200
num_images_per_prompt = 5

generator = torch.manual_seed(seed)
images = pipe(
    prompt = prompt,
    negative_prompt = negative_prompt,
    generator = generator,
    num_images_per_prompt = num_images_per_prompt).images

for i, image in enumerate(images):
    image.save(f"result_{i}.png")



赤い髪と赤い眼をネガティブプロンプトとした場合の結果をのせておきます。

ネガティブプロンプトなし
ネガティブプロンプトあり

実行例(年賀状のイラストを描いてみた)

来年は兎年です。

そして我が家は両親と子供の3人家族です。

年賀状のために3頭の兎のイラストを描いてみました。

prompt: An illustration of 3 rabbits
negative prompt: None
seed: 203
num_images_per_prompt: 5

prompt: An illustration of 3 rabbits,award-winning
negative prompt: None
seed: 203
num_images_per_prompt: 5

prompt: An illustration of 3 rabbits,award-winning,no background
negative prompt: None
seed: 203
num_images_per_prompt: 5

prompt: a low poly illustration of 3 rabbits,award-winning
negative prompt: None
seed: 203
num_images_per_prompt: 5

prompt: a low poly illustration of 3 rabbits,award-winning,white background
negative prompt: None
seed: 203
num_images_per_prompt: 5

prompt: a low poly illustration of 3 rabbits,award-winning,no background
negative prompt: None
seed: 203
num_images_per_prompt: 5