Stable Diffusionをpythonで動かしたい

はじめに

本記事では、Stable DiffusionをPythonのスクリプトで実行することを目的とする。

アカウントを作成する

アカウントを登録するときに、「[Hugging Face] Click this link to confirm your email address」という題のメールが届いているのでそちらを確認する必要がある。

トークン発行

[Settings]をクリックする。

[Access Tokens]をクリックする。

[New token]をクリックする。

[Name]に任意の文字列を記入する。

アクセストークンをコピーして控えておく。

仮想環境の構築

Pythonの独立した環境を作成したい場合は仮想環境を構築されたい。

仮想環境 venv を構築する。

python -m venv venv

仮想環境を有効にする

venv\Scripts\activate

マシンの構成確認

実行環境に CUDA ( NVIDIA 社のGPUプログラム開発環境) があるか否かでプログラムの一部が異なる。次のコマンドで確認せよ。

wmic path win32_videocontroller get caption

次に2つの実行例を示す。

実行例A: NVIDIA 製の GPU 搭載の場合

$ wmic path win32_videocontroller get caption
Caption
NVIDIA GeForce GTX 1080

実行例B: CPU 内蔵 GPU の場合

$ wmic path win32_videocontroller get caption
Caption
Intel(R) UHD Graphics 620

パッケージのインストール

まずは、深層学習のライブラリとしてPyTorchもしくはFlaxのうちどちらかをインストールする必要がある。この記事では、PyTorchの場合のみについて示す。

参考程度にFlaxの場合は次のコマンドでインストールされたい。

pip install diffusers["flax"] transformers

公式サイト にあった次のコマンドでインストールしていく。

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

次のコマンドで数値（バージョンを示す）が返されれば成功している。

python -c "import torch; print(torch.version.cuda)"

次に HuggingFace が開発しているライブラリ Diffusers をインストールする。深層学習のライブラリ（PyTorch/Flax）によってコマンドが違うので注意せよ。

• PyTorchの場合

pip install diffusers["torch"] transformers

プログラム

実行例

test.py

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline

access_token = トークン

# モデルのダウンロード
model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=access_token)
model.to("cpu")

# import matplotlib.pyplot as plt

num = 3 # 生成したい画像枚数
prompt = "Horse flying in the sky" # 生成させたいもととなる文章
 

# モデルの出力構造を確認
output_structure = model(prompt)

# 出力構造を表示
print(output_structure)

for i in range(num):
    # image = model(prompt)["sample"][0]
    image = output_structure.images[0]

    # 保存
    image.save(f"outputs/{i}.png")

参考