【Stable Diffusion】呪文・引数を画像に保存する

【Stable Diffusion】呪文・引数を画像に保存する 機械学習

「効果のあった呪文を忘れてしまった・・・」
「guidance_scaleやnum_inference_stepsの数値を保存しておきたい・・・」

このような場合に、この記事の内容が参考になります。
この記事では、効果のあった呪文・引数を画像に保存する方法を解説しています。

本記事の内容

  • なぜ、呪文・引数を保存したいのか?
  • 画像と呪文・引数を関連付ける方法
  • コードを用いた動作確認

それでは、上記に沿って解説していきます。

なぜ、呪文・引数を保存したいのか?

呪文とは、Stable Diffusionにおけるプロンプトのことです。
一般的には、これを呪文と言います。

そして、引数とは次のようなパラメータの値を指します。

  • guidance_scale(プロンプトと出力画像の類似度)
  • num_inference_steps(画像生成に費やすステップ数)
  • strength(入力画像と出力画像と相違度)
  • height(512がデフォルト)
  • width(512がデフォルト)

具体的には、pipeに渡す引数のことです。

with autocast("cuda"): 
    image = pipe(引数)["sample"][0]

それであれば、seedやプロンプトも引数と言えます。
みなさんは、これらの値を保存しておきたいと思ったことはありませんか?

おそらく、みなさんは魔法使いの見習いように何度も呪文を唱えることでしょう。
唱える呪文を途中で変更することは、何度もあるでしょう。

そして、微調整のために引数の値を変更することもあるはずです。
最終的に手元にそれらの値が残っていれば、ラッキーと言えます。

しかし、試行錯誤を重ねるうちに呪文や引数を失うことはあり得ます。
だからこそ、これらの値を保存したいと考えるのです。

試行錯誤に熱中すると、次のように画像だけが手元に残ります。

どの画像が、どんな呪文や引数で作ったのかはわかりません。
だからと言って、都度記録することは手間がかかります。

以上、呪文・引数を保存したい理由について説明しました。
次は、画像と呪文・引数を関連付ける方法を説明します。

画像と呪文・引数を関連付ける方法

呪文・引数といった画像生成の条件を保存したい理由を説明しました。
しかし、これらを単純に保存するだけでは意味がありません。

生成された画像と呪文・引数が、セットになって初めて意味があると言えます。
セットにする方法は、複数考えられます。

今回は、画像のExif情報を利用する方法を説明します。
なお、ここでの説明は次の記事の内容が済みであることを前提とします。

「txt2img.py」スクリプトをコマンドで実行する場合は、ここからの説明は無意味になります。
「txt2img.py」スクリプトを利用するケースは、システムとしての拡張性がありません。

あくまで、コマンドラインツールとしてStable Diffusionを利用するだけです。
拡張性を求める場合は、上記記事を参考にしてStable Diffusionをインストールしてください。

前提となる環境が整っているなら、次へ進みましょう。

以上、画像と呪文・引数を関連付ける方法を説明しました。
次は、コードを用いた動作確認を説明します。

コードを用いた動作確認

画像と呪文・引数を関連付けるためのコードは、以下。

import torch 
from diffusers import StableDiffusionPipeline 
from torch import autocast 
import datetime 
import json 
import pyexiv2 

MODEL_ID = "CompVis/stable-diffusion-v1-4" 
DEVICE = "cuda" 
YOUR_TOKEN = "コピーしたアクセストークン" 
PROMPT = "One Chihuahua-elephant hybrid" 
SEED = 1020 
STEP = 130 
SCALE = 7.5 

pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, revision="fp16", torch_dtype=torch.float16, 
                                               use_auth_token=YOUR_TOKEN) 
pipe.to(DEVICE) 

# seed固定 
generator = torch.Generator(device=DEVICE).manual_seed(SEED) 
# list型に引数・呪文を格納→json文字列に変換 
dic_info = {"prompt": PROMPT, "guidance_scale": SCALE, "seed": SEED, "num_inference_steps": STEP} 
json_str = json.dumps(dic_info) 

with autocast(DEVICE): 
    image = pipe(PROMPT, guidance_scale=SCALE, num_inference_steps=STEP, generator=generator)["sample"][0] 
    # 現在時間 
    dt_now = datetime.datetime.now() 
    now = dt_now.strftime("%Y%m%d%H%M%S") 
    # ファイル名 
    file_path = str(SEED) + "_" + str(now) + ".png" 
    # ファイル保存 
    image.save(file_path) 
    # Exif編集 
    with pyexiv2.Image(file_path) as img: 
        img.modify_exif({'Exif.Image.ImageDescription': json_str})

コードのベースとなる部分は、次の記事で説明しています。

seedについては、次の参考にしてください。

そして、Exif情報への書き込みについては次の記事で解説しています。
コードで言うと、次の部分です。

    # Exif編集 
    with pyexiv2.Image(file_path) as img: 
        img.modify_exif({'Exif.Image.ImageDescription': json_str})

あとは、実際に動かしてみればわかるでしょう。
実行すると、1020_で始まる画像が作成されます。
1020_以降は、処理日時です。

私の環境では、次の画像が作成されます。
おそらく、seedが同じであれば同じ画像が作成されるはずです。

この画像には、Exifが書き込まれています。
Exifは、次のようなWebツールで確認できます。

写真・画像のexif/メタデータ等情報確認・削除【無料】
https://image-convert.cman.jp/imgInfo/

このようなWebツールを使わずに、Pythonで読み取りも可能です。
pyexiv2は、Exifの確認でも利用できます。

import json 
import pyexiv2 
import pprint 

IMG_PATH = "1020_で始まるファイル.png" 

with pyexiv2.Image(IMG_PATH) as img: 
    data = img.read_exif() 
    json_str = data['Exif.Image.ImageDescription'] 
    dic_data = json.loads(json_str) 
    pprint.pprint(dic_data)

上記コードを実行すると、次の結果が表示されます。

{'guidance_scale': 7.5, 
 'num_inference_steps': 130, 
 'prompt': 'One Chihuahua-elephant hybrid', 
 'seed': 1020}

json文字列としてExifに書き込んでいます。
そのため、プログラムからは利用しやすいです。

以上、コードを用いた動作確認を説明しました。

タイトルとURLをコピーしました