「Stable Diffusionで作成される画像の精度を少しでも上げたい」
「Stable DiffusionにおけるVAEを変更したい」
このような場合には、この記事の内容が参考になります。
この記事では、Stable DiffusionでVAEを変更して画質を上げる方法を解説しています。
本記事の内容
- VAEとは？
- Stable DiffusionにおけるVAE
- Stable DiffusionでVAEを指定する方法
それでは、上記に沿って解説していきます。
VAEとは？
VAEを説明する前に、AutoEncoder（オートエンコーダ）から説明しておきます。
AutoEncoderは、ニューラルネットワークを利用した教師なし機械学習の手法の一つです。
少しだけ難しいAutoEncoderの説明は、以下となります。
入力されたデータを一度圧縮し、重要な特徴量だけを残した後、再度もとの次元に復元処理をするアルゴリズム
そして、VAEはVariational AutoEncoderの略称になります。
AutoEncoderとVAEは関係があるのがわかりますね。
では、その違いは何になるのでしょうか？
簡単に言うと、AutoEncoderの潜在変数に確率分布を用いたのがVAEになります。
「確率分布って何やねん！」と思うかもしれません。
確率分布を使えば、予測の精度が高まるとだけ覚えておきましょう。
つまり、VAEを利用するとAutoEncoderよりも予測の精度が高まるということです。
もう少し言うと、未知に対する対応が取りやすくなるとも言えます。
以上、VAEについて説明しました。
次は、Stable DiffusionにおけるVAEを説明します。
Stable DiffusionにおけるVAE
Stable DiffusionにおいてもVAEが利用されています。
オリジナルなVAEであり、「kl-f8 autoencoder」と名前が付けられています。
デフォルトでは、「kl-f8 autoencoder」がVAEとして利用されているということです。
そんな状況において、新たにVAEが2つ公開されました。
- ft-EMA
- ft-MSE
VAEが異なることによる効果は、次のようになると紹介されています。
256×256: ft-EMA (left), ft-MSE (middle), kl-f8 autoencoder(right)
目や文字が潰れなくなっていますね。
このような効果を簡単に出せるなら、是非とも新たなVAEを利用したいとなるはずです。
次に、3つのVAEの学習ステップ数の違いを確認しましょう。
| Model | train steps |
|---|---|
| kl-f8 autoencoder | 246803 |
| ft-EMA | 560001 |
| ft-MSE | 840001 |
ft-MSEが、最も学習しているということですね。
ただし、これはデコーダ部分（復元処理）における話です。
デコーダ部分のみの調整であるため、影響範囲を限定できます。
限定できるからこそ、既存モデルをそのまま利用可能となります。
「CompVis/stable-diffusion-v1-4」を別のモデルに置き換える必要がないということです。
また、各VAEの評価指数が公開されています。
COCO 2017 (256×256, val, 5000 images)
| Model | rFID | PSNR | SSIM | PSIM |
|---|---|---|---|---|
| kl-f8 autoencoder | 4.99 | 23.4 +/- 3.8 | 0.69 +/- 0.14 | 1.01 +/- 0.28 |
| ft-EMA | 4.42 | 23.8 +/- 3.9 | 0.69 +/- 0.13 | 0.96 +/- 0.27 |
| ft-MSE | 4.70 | 24.5 +/- 3.7 | 0.71 +/- 0.13 | 0.92 +/- 0.27 |
LAION-Aesthetics 5+ (256×256, subset, 10000 images)
| Model | rFID | PSNR | SSIM | PSIM |
|---|---|---|---|---|
| kl-f8 autoencoder | 2.61 | 26.0 +/- 4.4 | 0.81 +/- 0.12 | 0.75 +/- 0.36 |
| ft-EMA | 1.77 | 26.7 +/- 4.8 | 0.82 +/- 0.12 | 0.67 +/- 0.34 |
| ft-MSE | 1.88 | 27.3 +/- 4.7 | 0.83 +/- 0.11 | 0.65 +/- 0.34 |
画質の評価については、次のことがわかっています。
PSNRが高ければ高いほど、画質が良い SSIMが高ければ高いほど、画質が良い
それ以外の項目は、よくわかりません。
でも、傾向を見るとrFIDもPSIMも低いほど良さそうです。
とにかく、ft-MSEの結果が最も良いということはわかります。
やはり、ステップ数が多いほど良いのですね。
以上、Stable DiffusionにおけるVAEを説明しました。
次は、Stable DiffusionでVAEを指定する方法を説明します。
Stable DiffusionでVAEを指定する方法
stabilityai/sd-vae-ft-ema · Hugging Face
https://huggingface.co/stabilityai/sd-vae-ft-ema
公式では、次のコードが掲載されています。
from diffusers.models import AutoencoderKL
from diffusers import StableDiffusionPipeline
model = "CompVis/stable-diffusion-v1-4"
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema")
pipe = StableDiffusionPipeline.from_pretrained(model, vae=vae)
通常は、このままでは動きません。
あくまで目安にするぐらいです。
サンプルをベースに改良したのが、次のコードになります。
ft-EMAを利用する場合の動作確認用です。
from diffusers import StableDiffusionPipeline
import torch
from diffusers.models import AutoencoderKL
MODEL_ID = "CompVis/stable-diffusion-v1-4"
DEVICE = "cuda"
YOUR_TOKEN = "コピーしたアクセストークン"
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema")
pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float16, revision="fp16",
use_auth_token=YOUR_TOKEN, vae=vae)
pipe = pipe.to(DEVICE)
prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]
image.save("astronaut_rides_horse.png")
「コピーしたアクセストークン」については、次の記事内で説明しています。
上記記事の「Access Token（アクセストークン）の取得」という箇所をご覧ください。
上記コードは、ft-EMAの場合でした。
ft-MSEの場合は、「ema」を「mse」に置き換えるだけです。
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
さて、コードを実行すると、初回時にはチェックポイントをダウンロードします。
そのため、処理に少し待たされるかもしれません。
もちろん、VAEを置き換えていることで処理が遅くなっていることもあるでしょう。
そして、作成できた画像は次の画像です。
正直、VAEを置き換えた効果はわかりません。
seedを固定して、各VAEを試してみればわかるのかもしれませんね。
以上、Stable DiffusionでVAEを指定する方法を説明しました。
