Stable Diffusionを試してみる
Stable Diffusion は、CompVis、Stability AI、LAION の研究者とエンジニアによって作成された、テキスト文章から画像を生成するモデルです。LAION-5Bデータセットでトレーニングされたモデルを提供しています。
Stable Diffusionのモデルはこちら
https://huggingface.co/CompVis/stable-diffusion
デモはこちら
https://huggingface.co/spaces/stabilityai/stable-diffusion
WindowsのDocker環境で動かしてみる
今回もWindowsのdocker環境を使います。GPUが使えるのは嬉しいですね！環境のセットアップに関しては過去記事を参考にしてください。
早速、pytorchのコンテナ使います。WindowsのDドライブに画像を出力したいので、コンテナにマウントしています。
docker run --gpus all -it -v D:\work:/work nvcr.io/nvidia/pytorch:22.08-py3
コンテナが起動したらパッケージをインストールします。
pip install diffusers==0.2.4 transformers scipy ftfy
あとは、簡単なプログラムを書いて動かしてみます。
import torch
import random
from diffusers import StableDiffusionPipeline
USER_ACCESS_TOKEN="<Hugging FaceのSettingsから作成するUser Access Token>"
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=USER_ACCESS_TOKEN).to("cuda")
prompt = "A photo of a raccoon wearing an astronaut helmet, looking out of the window at night."
with autocast("cuda"):
image = pipe(prompt)["sample"][0]
image.save(f"image.png")
Stable Diffusionは、Imagen や、Midjourney のようにテキスト文章を入力して画像を生成します。プログラムでは、テキスト文章をpromptに設定して画像を生成しています。今回はImagenのページに紹介されている同じテキスト文章を入力して試しました。このプログラムは実行されるたびに異なる画像が出力されます。なお、ユーザーアクセストークンは、Hugging Faceにログインした後に、右上のアイコンからSettingsを選択し、User Access Tokenを取得します。
あと、GPUのメモリが確保できず動作しない場合は、FP16でモデルで動かすことにより使用するメモリのサイズを減らして動作させることができます。FP16だと5GB～6GBを使います。以下のように、StableDiffusionPipeline.from_pretrainedの箇所でFP16のモデルを選択してください。
import torch
from torch import autocast
import random
from diffusers import StableDiffusionPipeline
USER_ACCESS_TOKEN="<Hugging FaceのSettingsから作成するUser Access Token>"
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4",
torch_dtype=torch.float16,
revision="fp16",
use_auth_token=USER_ACCESS_TOKEN).to("cuda")
prompt = "A photo of a raccoon wearing an astronaut helmet, looking out of the window at night."
with autocast("cuda"):
image = pipe(prompt)["sample"][0]
image.save(f"image.png")
自分のイメージしている絵を探す
自分のイメージしている絵を探すために沢山の画像を作ってみます。もしかしたら、ノイズがひどいかもしれませんので後からノイズの少ない同じ画像を作れるように、Seed値を把握できるようにしておきます。
Stable Diffusion の画像生成は、ユーザのプロンプト（テキスト文章）と、Seed値（乱数）を入力として画像が生成されます。テキスト文章はCLIPのエンコーダーを経由して、Seed値と共にUNetに入り、最後はVAのデコーダーで出力されるようですので、テキスト文章とSeed値を固定にすることにより同じ画像を作ることができます。
import torch
from torch import autocast
import random
from diffusers import StableDiffusionPipeline
USER_ACCESS_TOKEN="<Hugging FaceのSettingsから作成するUser Access Token>"
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=USER_ACCESS_TOKEN).to("cuda")
prompt = "A photo of a raccoon wearing an astronaut helmet, looking out of the window at night."
for i in range(100):
seed = random.randrange(0, 2147483647, 1)
with autocast("cuda"):
generator = torch.Generator("cuda").manual_seed(seed)
image = pipe(prompt, generator=generator)["sample"][0]
image.save(f"astronaut_seed_{seed}.png")
ファイル名にSeed値を入れているので、気に入った画像があればSeed値からノイズを減らした画像を後から生成することができます。黒い画像がある場合は、安全チェックににより NSFW と判断されたものになります。
画像のノイズを減らす
気に入った画像があれば、Seed値を指定してnum_inference_stepsを大きくすることにより、よりノイズの少ない画像を生成することができます。
import torch
from torch import autocast
import random
from diffusers import StableDiffusionPipeline
USER_ACCESS_TOKEN="<Hugging FaceのSettingsから作成するUser Access Token>"
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=USER_ACCESS_TOKEN).to("cuda")
seed = 1160424331
prompt = "A photo of a raccoon wearing an astronaut helmet, looking out of the window at night."
with autocast("cuda"):
generator = torch.Generator("cuda").manual_seed(seed)
image = pipe(prompt, guidance_scale=7.5, height=512, width=512, num_inference_steps=100, generator=generator)["sample"][0]
image.save(f"astronaut_{seed}_step_100.png")
num_inference_stepsの回数により画像がどのように変化するか確認してみました。右上から左下にかけてnum_inference_stepsの値を、5～100まで変化させています。
FP16とFP32で画質の違いはあるの？
以下はFP16で生成した画像です。上記のFP32と同じ入力になります。（FP32はRTX3090、FP16はRTX2070を使用して生成しています）
微妙に違いはありますが、ほぼ変わらないですね。
さいごに
Stable Diffusion 凄いですね！思いもよらない構図の画像が出てきたときは驚きます。以下はStable Diffusionが生成した、東京に大きな穴がある画像です。凄くありそうですし、ビルと穴の位置に違和感がありません。凄い。。。
