本記事では、DreamBoothと呼ばれる機械学習手法を用いて事前学習済みStable Diffusionを追加学習する方法をご紹介します。

Dream Booth

概要

Dreamboothは、Text to Imageタスクを実現する拡散モデルの追加学習手法です。

Dreamboothでは、まず、事前に学習されたText to Imageの拡散モデルに数枚の特定の被写体が写る画像と、識別子となるプロンプトを与え、追加学習させます。

追加学習されたモデルは、識別子を用いた様々なプロンプトに応じて画像を生成することが可能となります。

上図では、いらすとや様の男の子の画像を数枚追加学習させ、「自転車に乗る男の子」をモデルに出力してもらった結果を示しています。

追加学習させたモデルでは、識別子「男の子」を与えると、いらすとや様の男の子風の画像を生成することが可能になっています。

このように、数枚の画像を追加学習させることにより、モデルをカスタマイズすることが可能になります。

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、Stable Diffusionの事前学習済みモデルを追加学習していきます。

デモ(Colaboratory)

それでは、実際に動かしながらStable Diffusionの追加学習を行っていきます。
本記事にはソースコードの要点を記載しています。全文は下記のGitHubをご参照下さい。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきますが、Hugging FaceのAccess Tokenは取得済みであるという前提のもと進めていきます。
Access Tokenの取得がまだの方は以下の記事をご参照ください。

Access Tokenの取得後、Colaboratoryを開き、下記を設定しGPUを使用するようにしてください。

はじめに、ライブラリをインストールします。

!pip install -qq git+https://github.com/huggingface/diffusers.git accelerate tensorboard transformers ftfy gradio
!pip install -qq "ipywidgets>=7,<8"
!pip install -qq bitsandbytes

次に、Hugging Faceへのログインを済ませておきます。
以下のセルを実行後表示されるプロンプトに取得したAccessTokenを入力します。

from huggingface_hub import notebook_login
!git config --global credential.helper store

notebook_login()

最後にライブラリをインポートします。

import argparse
import itertools
import math
import os
from contextlib import nullcontext
import random

import numpy as np
import torch
import torch.nn.functional as F
import torch.utils.checkpoint
from torch.utils.data import Dataset

import PIL
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from diffusers import AutoencoderKL, DDPMScheduler, PNDMScheduler, StableDiffusionPipeline, UNet2DConditionModel
from diffusers.optimization import get_scheduler
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
from PIL import Image
from torchvision import transforms
from tqdm.auto import tqdm
from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer

import bitsandbytes as bnb

以上で環境セットアップは完了です。

学習済みモデルのセットアップ、追加学習画像の設定

ここでは、追加学習に用いるベースとなるモデルを設定します。

今回はStable diffusionを使用します。

pretrained_model_name_or_path = "CompVis/stable-diffusion-v1-4" #@param {type:"string"}

続いて、追加学習させる画像を指定します。

ここでは、いらすとや様の男の子の画像を5枚使用します。

urls = [
    'https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEiCSk6WI_XdfyJdIab2Iebe3KZtcvlN4izAAJZMENM2qz7UH7QP3_M-zp6bmopKLTGo4_Z6M1sd0-ZRm7xFQeH1xbCjteSL4EBdygOiEfoUGP5Gx7QHrvZSxLusmgcQQerQofQHNMEts64y/s250/penlight_man01_blue.png',
    'https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEgZ5F3JkOSY-AsniE91XOcEEVf_hpTrc07G9VdOo0AuqEIOSFLDb96OScI00vZsQfcVfzM5nwDmRg3m40LNv0W10Gz0vpTa8WOMn5wUJMk6leURo9Owb76VnmqrkFx-2dcJjGv_K2IeobER/s250/penlight_man05_orange.png',
    'https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEirusu8jzE5W0PVsa9Czgx5izBJIpKHomRYROtGKmly0_WSpFPtoJdskjTW3nW-4BEaVkM7Mm_XjCTfA8tbXCw-TRPVENIthTmqWbBDx_UXVkdOHeykUIKNz7MrVX8-6O8ZYEZwhYoyxM9C/s330/stand_naname1_boy.png',
    'https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEjdJK13Y1oP76dyVbVtPCI6tq1i1EfOzzRNrElxEvXB7lFCCr3GJphn1yCtMfqTugRa24LsfqRcylAbwjU5s501QbG1C4mXY4EgW3QxLCO8xf5lsUCCejoU_au3sxZhsYg62JlnwgIJ2gsK/s380/stand_naname2_school_boy.png',
    'https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEg66dr3gmrwA9X9hRQYJgZKMUh5E6xII5blwlOmZT_t3jYDTSoAfkRJ_gYnn0ChnDpC1qi8svE2X99cU9SMq1QSj1CIzjj7IlPiv9hza05QT7HKoUjR33P4qU7kv5sKQG_jgGjFzGwdewSQ/s400/stand_naname3_man.png',
    'https://blogger.googleusercontent.com/img/b/R29vZ2xl/AVvXsEgtJcfF9UV1libQ1VhBXGxqt98VL1vObciKA6HUsUpD9jM7IcFncOBuqrbDZa2uoN1NMpy-11WGxlg9G_zlHrVXQucRB8mGAbmdDHnJjGh8CvXS1axe9PTyfaziuQhcVpqb9XFF1qknJqTC/s400/stand_naname4_businessman.png'
          ]

識別子の設定

ここでは、学習させる被写体の識別子を定義します。

追加学習させた識別子であることを示す"sks"+オブジェクト名でプロンプトを指定します。
今回はan illustration of sks boyを指定します。

# オブジェクト、画風の説明
instance_prompt = "an illustration of sks boy" #@param {type:"string"}

# コンセプトクラスの指定、画質が向上
prior_preservation = False #@param {type:"boolean"}
prior_preservation_class_prompt = " an illustration of a boy" #@param {type:"string"}

num_class_images = 12 
sample_batch_size = 2
prior_loss_weight = 0.5
prior_preservation_class_folder = "./class_images"
class_data_root=prior_preservation_class_folder
class_prompt=prior_preservation_class_prompt

あとは残りのセルをすべて実行することでトレーニング始まります。
無料のGoogle Colaboratory環境(Tesla T4)で約30分ほどで完了します。

追加学習モデルによるText to Image

それでは、追加学習させたモデルで様々な画像を生成していきます。

pipe = StableDiffusionPipeline.from_pretrained(
        args.output_dir,
        torch_dtype=torch.float16,
    ).to("cuda")

from torch import autocast
prompt = "an real illustration of sks boy swimming sea" #@param {type:"string"}

num_samples = 1 #@param {type:"number"}
num_rows = 1 #@param {type:"number"}

all_images = [] 
for _ in range(num_rows):
    with autocast("cuda"):
        images = pipe([prompt] * num_samples, num_inference_steps=50, guidance_scale=7.5).images
        all_images.extend(images)

grid = image_grid(all_images, num_samples, num_rows)
grid

プロンプトan real illustration of sks boy swimming seaの出力結果は以下の通りです。

an illustration of sks boy eating hamburgerは以下の通りです。

いらすとや様の男の子の画風を継承した様々な画像が生成されています。

まとめ

本記事では、DreamBoothを用いたStable Diffusionの追加学習方法をご紹介しました。

高い表現力を誇るStable Diffusionを使って特定のオブジェクトに特化した画像が生成できるので求める画像が生成しやすくなります。
一方で、悪用厳禁であることは言うまでもありません。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation

2. GitHub - XavierXiao/Dreambooth-Stable-Diffusion

3. GitHub - XavierXiao/Dreambooth-Stable-Diffusion