Stable Diffusionをローカル環境で使用する方法

この記事は公開から1年以上経過しています。

私がアイキャッチやサムネイル作成などに利用しているStable Diffusionのローカル環境(最小限のコード)を紹介します。

一般的にはDockerなどを使ってWEBフロントエンドを使えるようにするようですが、自分で即席であれこれ弄れるほうが楽しいので…。

ちなみにChatGPT APIのようなLLMと連携させることで、自然言語指示から画像生成プロンプトを生成して画像を生成するとか、プロンプトで適切なモデルに切り替える…のような応用(改造)も可能だと思います。

環境はLinux(Ubuntu 22)＋NVIDIA RTX3060(12GB)＋Python(3.10.12)です。
GPUなしでも実行できますが、RTX3060で1枚約8秒のところCPU(i7 13700K)では約150秒位掛かりました。
(空冷CPUのためサーマルスロットリングが掛かっている可能性大。)

本エントリで紹介する方法は画面から設定するなどの機能はありません。あくまでPythonプログラミングが多少分かる方向けの解説となりますので予めご了承ください。

サンプルソースコード

pip:

accelerate
diffusers
huggingface-hub
omegaconf
torch
transformers

Hugging Faceモデル使用版:

use_huggingface_hub.py

import time

import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import login
from torch.amp.autocast_mode import autocast

DEVICE = 'cuda'

# モデルダウンロード先ディレクトリ
CACHE_DIR = 'cache/diffusers'

# Hugging Face Hubにログイン
login(token='Hugging Face Token Id')

# モデルリポジトリ名
model_id = 'CompVis/stable-diffusion-v1-4'

if DEVICE == 'cuda':
    # VRAMフラグメント対策
    torch.cuda.memory._set_allocator_settings('max_split_size_mb:100')

pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    cache_dir=CACHE_DIR,
    safety_checker=None)
pipe = pipe.to(DEVICE)

prompt = """
ここにプロンプトを記載
"""

negative_prompt = """
ここにネガティブプロンプトを記載
"""

# ループ処理
for i in range(0, 10):
    with autocast('cuda'):
        image = pipe(prompt=prompt, negative_prompt=negative_prompt, guidance_scale=7.5).images[0]
    # タイムスタンプ付きファイル名で保存
    fileName = f'img-{time.strftime(r"%Y%m%d%H%M%S")}.png'
    image.save(fileName)

解説

モデルのダウンロードサイズが数GBあります。通信環境にご注意ください。

DEVICEに、GPUを使用するかCPUを使用するかを指定します。
cudaを指定するとGPU、cpuを指定するとCPUで処理を行います。
CACHE_DIRに、モデルダウンロード先パスを入力します。
この例ではスクリプトと同じディレクトリ内の/cache/diffusersディレクトリを設定しています。
Hugging Face Token Idに、ご自身のHuggingFaceのトークンIDを設定します。
Hugging FaceのトークンIDについてはこちらで解説しています。
model_idに、画像生成モデル名を入力します。
この例ではこちらのモデルが設定されています。
safety_checkerに、NSFW(Not Safe For Work＝職場回覧注意)の解除を指定しています。
オフィスなどで利用する場合はこのパラメータを削除しておくことをオススメします。
promptにプロンプトを入力します。
negative_promptにネガティブプロンプトを入力します。

このスクリプトを実行すると、指定されたモデルとプロンプトで10枚のPNG画像をスクリプトと同じディレクトリ内にimg-YYYYMMDDHHMMSS.pngというファイル名で保存します。

ローカルモデル(safetensors)使用版:

use_safetensors.py

import time

import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import login
from torch.amp.autocast_mode import autocast

DEVICE = 'cuda'

# ローカルモデルパス(safetensors)
model_path = 'path/to/model.safetensors'

if DEVICE == 'cuda':
    # VRAMフラグメント対策
    torch.cuda.memory._set_allocator_settings('max_split_size_mb:100')

pipe = StableDiffusionPipeline.from_single_file(model_path, load_safety_checker=False)
pipe = pipe.to(DEVICE)

prompt = """
ここにプロンプトを記載
"""

negative_prompt = """
ここにネガティブプロンプトを記載
"""

# ループ処理
for i in range(0, 10):
    with autocast('cuda'):
        image = pipe(prompt=prompt, negative_prompt=negative_prompt, guidance_scale=7.5).images[0]
    # タイムスタンプ付きファイル名で保存
    fileName = f'img-{time.strftime(r"%Y%m%d%H%M%S")}.png'
    image.save(fileName)

解説

DEVICEに、GPUを使用するかCPUを使用するかを指定します。
cudaを指定するとGPU、cpuを指定するとCPUで処理を行います。
model_pathに、ローカルにダウンロード済のsafetensorsモデルのパスを設定します。
CIVITAIなどからダウンロードできるモデルを使用できます。
load_safety_checkerに、NSFW(Not Safe For Work＝職場回覧注意)の解除を指定しています。
オフィスなどで利用する場合はこのパラメータを削除しておくことをオススメします。
promptにプロンプトを入力します。
negative_promptにネガティブプロンプトを入力します。

デフォルトの画像サイズは512×512です。変更したい場合は、pipe()にwidthとheightパラメータを追加してください。その他利用可能なパラメータについては、参考ウェブサイトをご参照ください。

参考ウェブサイトなど

Hugging Face Docs
Text-to-image/StableDiffusionPipeline

以上です。