AIに自分のペットの絵を描かせてみた（StableDiffsuion, dreambooth）

こんにちは、すきにーです。

今日は流行りの画像生成AI,　StableDiffsuionで飼い猫であるちゃおの絵を描いてみました。

元の写真

AIに描かせた猫

すごくないですか？！？！？
ここまでそっくりになるとは思わなかったです。
画像は荒いけど感動しました。

どうやったらこのようなそっくりの絵を描けるのか。
今日はその方法を書いていきたいと思います。

①StableDiffusionとは？
Stable Diffusion は、潜在的なテキストから画像への拡散モデルであり、任意のテキスト入力が与えられた場合に写真のようにリアルな画像を生成し、信じられないほどの画像を生成するための自律的な自由を育成し、何十億もの人々が数秒で見事なアートを作成できるようにします。（公式サイトから）

例えば可愛い女の子の画像を描いてと言ったらこのような画像が出てきます。
（実際はもっと詳細な命令が必要ですが）

これが無料で利用でき、商用利用もOKなので凄いです。

しかし、StableDiffusionを使うだけでは自分のペットの画像を出力することはできません。
何故ならStableDiffusionは、全般的な何億枚もの画像を学習していて、
自分のペットに特化させて学習させているわけではないからです。

そこでdreamboothを使います。

②DreamBoothとは？
DreamBoothとは、特定の対象を事後学習させる技術です。
StableDiffusionの全般的な画像学習に加えて、DreamBoothで自分のペットの画像を「自分のペット」とラベル付けすることで、自分のペットの画像を生成できるAIモデルを作ることができます。
(【個人開発】AIに「自分のお気に入り」の画像を作らせよう【DreamBooth】より)

③作り方
ここではdreamboothで自分のペットを作る方法の一部を紹介します。

dreamboothのコードをgithubからダウンロード

!wget -q https://github.com/ShivamShrirao/diffusers/raw/main/examples/dreambooth/train_dreambooth.py

{key}に自分のペット
{className}にペットの種類のフォルダを作り、画像を格納します。
classNameは猫ならcat

concepts_listでコンセプトを作ります

concepts_list = [
    {
        "instance_prompt":      f"photo of {key} {className}",
        "class_prompt":         f"photo of a {className}",
        "instance_data_dir":    f"/content/data/{key}",
        "class_data_dir":       f"/content/data/{className}"
    }
]

ipynb.ipynb でダウンロードしたtrain_dreambooth.py を実行します。

ipynb.ipynb

!accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base" \
  --pretrained_vae_name_or_path="stabilityai/sd-vae-ft-mse" \
  --revision="fp16" \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --seed=1337 \
  --resolution=512 \
  --train_batch_size=1 \
  --train_text_encoder \
  --mixed_precision="fp16" \
  --use_8bit_adam \
  --gradient_accumulation_steps=1 \
  --learning_rate=1e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --sample_batch_size=1 \
  --max_train_steps=1600 \
  --save_interval=100000 \
  --concepts_list="concepts_list.json"

体感ですが、max_train_stepsは1200から2400がよさげです。
これで学習が完了です。

次にpromptを指定します。

f"portrait of chao cat, vibrant dynamic portrait by makoto shinkai",

指定したpromptで画像を生成します。

pipe = StableDiffusionPipeline.from_pretrained(model_path, scheduler=scheduler, safety_checker=None, torch_dtype=torch.float16).to("cuda")
images = pipe(
  prompt,
  num_images_per_prompt=3,
  generator=g_cuda
).images
print(images)

ペットが描かれます。

⑤終わりに
AIの画像生成がここまで進化してるとは思いませんでした。
初めてAIで画像を作ってみたけどとても楽しかったですね。

いつか自分で生成モデルを作って絵を生成してみたいです。