この記事は，ドコモアドベントカレンダー17日目の記事になります。

こんにちは。ドコモの下山と言います。業務ではバーチャル空間における3Dコンテンツ生成に取り組んでいます。

この記事では，Stable Diffusionに代表されるText-to-Imageの技術と，Novel View Syntyesisの手法であるNeRF（Neural Radiance Fields）を組み合わせた手法であるDreamFusionを使って，手元に用意した10数枚のいらすとや画像から，いらすとや風の3Dモデルの生成を試みます。

DreamFusion

DreamFusionは"DreamFusion: Text-to-3D using 2D Diffusion"というタイトルの論文で今年の9月に発表された手法です。

現在，Stable Diffusionを始めとするText-to-Image手法が大きく盛り上がりを見せています。これらの手法は何億枚という大量の画像で学習することで実現されています。一方，これに匹敵する規模の3Dのデータセットは存在せず，これらの手法をそのまま3Dに拡張することは現実的ではありません。

このような課題意識で発表されたDreamFusionでは，学習に3Dデータを一切必要としません。ざっくりと手法を説明すると，事前学習済みのText-to-Imageモデルを用いて，NeRFを最適化するというものになっています。

いらすとや風の3Dモデルを作ってみる

今回はこのDreamFusionを使って，いらすとや風の3Dモデルを作ってみようと思います。

DreamFusionを使っていらすとや風の3Dモデルを作るうえで問題になるのは，いらすとや風の画像を生成するText-to-Imageモデルが存在しないことです。そこで，今回はDreamBoothという手法を使ってStable Diffusionをfine-tuningすることで，いらすとや風の画像を生成するText-to-Imageモデルを用意するところから始めたいと思います。

全体の大まかな流れは以下のようになります。

1) いらすとやの画像12枚と，公開されている事前学習済みのStable Diffusionを用意し，DreamBoothでfine-tuningします。 2) fine-tuningされたStable Diffusionと，任意のテキストから，DreamFusionにより3Dモデルを生成します。

DreamBoothによるStable Diffusionのfine-tuning

まず初めにDreamBoothについて簡単にご紹介します。 DreamBoothは"DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation"というタイトルの論文で今年の8月に発表された手法です。

DreamBoothではある特定の対象物を写した数枚の画像からfine-tuningをすることで，その対象に対応する識別子を新たに獲得します。推論時には獲得した識別子をテキストに含めることで，学習に用いた対象物が登場する新たな画像を自由に生成することが可能になります。

実装に関しては，Hugging FaceのDiffusersというパッケージでDreamBoothの実装が利用できるのですが，Diffusersをベースにこちらの解説記事の筆者が少し書き換えたものが分かりやすかったため，今回はこちらを使用します。

論文では3~5枚の画像で学習されていますが，上記の解説記事では12枚程度が良いと書かれていましたので，今回は12枚の画像を利用することにしています。

環境構築

はじめに手元の実行環境について簡単にご紹介します。

OS：Windows 11 Pro
GPU：GeForce RTX 3080 Ti Laptop
CUDA：11.7

今回使用するパッケージの中にWindowsでは動かないものがあるため，今回はWSL2のUbuntu 20.04 LTSディストリビューションを用いて検証を行います。

まずは今回用いる実装のREADME通りに必要なパッケージ等をインストールします。

$ pip install git+https://github.com/ShivamShrirao/diffusers.git
$ pip install -U -r requirements.txt

READMEにはこれしか記載されていないのですが，手元の環境ではこれだけだと動作しなかったため，ここから色々と手を加えていきます。

まず，インストールされたPyTorchがCPU版のものになっていたため，CUDA対応のものを入れなおします。

# CUDA 11.7対応のPyTorch
$ pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117

またこの実装ではメモリ効率を高めるため，8bitのAdamを使用しています。そのためbitsandbytesをインストールする必要があります（なぜかrequirements.txtに載っていない）。

# CUDA 11.7に対応したもの
$ pip install bitsandbytes-cuda117

このバージョンのbitsandbytesは非推奨なようでpip install bitsandbytesとするように警告が出るのですが，このやり方だと手元の環境ではbitsandbytesがGPUを認識しなかったため，今回は非推奨のバージョンを使用します。

また，これらの作業の過程で最新版のmarkupsafeがインストールされてしまうのですが，最新版ではエラーが出たため，少し古いものにダウングレードします。

$ pip install markupsafe==2.0.1

動作させる前の準備

READMEにある通り，accelerateの設定と，Hugging Faceのアクセストークンの設定を済ませる必要があります。まずaccelerateの設定を行います。

# accelerateの設定
$ accelerate config

すると5つの質問が準に表示されるので，それぞれ手元の実行環境に合わせて回答します。今回の環境では以下のように回答しました。

In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)) → 0
Which type of machine are you using? ([0] No distributed machine, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS) → 0
Do you want to run your training on CPU only (even if a GPU is available)? [yes/NO] → no
Do you want to use DeepSpeed? [yes/No] → no
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16] → fp16

続いて，Hugging Faceのアクセストークンを設定します。 WebブラウザでHugging Faceにログインし，settings > Access Tokensと進むとアクセストークンの管理画面にたどり着きます。ここでread権限のアクセストークンを作成します。アクセストークンが作成できたらコマンドライン上で設定を行います。

$ huggingface-cli login

トークンの入力を求められるので先ほど作成したアクセストークンをコピペすれば，設定は完了です。

動作確認

手元の環境ではGPUのVRAMが16GBなので，README記載の学習スクリプトの中からTraining on a 16GB GPUを選択し，編集します。

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export INSTANCE_DIR="path-to-instance-images" # 用意したデータを格納したディレクトリ
export CLASS_DIR="path-to-class-images" # 学習の際に生成される画像を格納するディレクトリ
export OUTPUT_DIR="path-to-save-model"

# 用意したデータに合わせてinstance_promptとclass_promptを編集する
accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --class_data_dir=$CLASS_DIR \
  --output_dir=$OUTPUT_DIR \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --instance_prompt="an illustration of sks girl" \ # sksは新たに追加する識別子
  --class_prompt="an illustration of girl" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=2 --gradient_checkpointing \
  --use_8bit_adam \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=200 \
  --max_train_steps=800

$INSTANCE_DIRには用意した画像（今回でいうと，いらすとやの画像）を格納します。今回の実験ではいらすとやの女の子の画像を用意したので，instance_promptは"an illustration of sks girl"としています。"sks"はfine-tuningにより新たに獲得させる識別子であり，ユニークな単語なら何にしても良さそうです。

また，DreamBoothでは過学習を抑えるために，prior-preservationというテクニックを使います。これは，instance_promptから識別子を取り除いたclass_prompt（今回でいうと"an iilustration of girl"）を使って，fine-tuning前にStable Diffusionに生成させた画像を学習に用いるというものです。class_promptから生成された画像は$CLASS_DIRに保存されます。

これを実行すると，学習が始まり，完了までに30分ほどかかります。学習が終わったら，いくつかのプロンプトから画像を生成し，fine-tuningが上手くできているか確認します。

学習したモデルから画像を生成するには次のコードを実行します。

from diffusers import StableDiffusionPipeline
import torch

model_id = "path-to-saved-model"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

prompt = "an illustration of sks girl"
image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]

image.save("path-to-dir-to-save-image")

promptの部分を書き換えることで様々な画像を生成できます。実際に実行してみた例をいくつか紹介します。

prompt	image
an illustration of sks girl
an illustration of sks girl eating a hamburger
an illustration of sks girl drinking a coffee
an illustration of sks girl riding a horse

最初の二つは，目などいらすとやの特徴を捉えているようにも見えますが，後の二つは...。特に最後のものは全く違った雰囲気のイラストになってしまっています。人の顔よりも馬の方が要素として大きく，馬の絵の画風の影響を強く受けてしまった可能性があります。

DreamFusionによる3Dモデル生成

DreamFusionは著者実装が公開されていません。そこで現在利用できる有志による実装を今回は利用します。

環境構築

上記の実装のREADMEに従って必要なものをインストールしていきます。

$ git clone https://github.com/ashawkey/stable-dreamfusion.git
$ cd stable-dreamfusion

# 必要なパッケージのインストール
$ pip install -r requirements.txt

# instant-ngpをバックボーンにしたモデルを動かすにはextensionのビルドが必要
$ bash scripts/install_ext.sh

動作確認

DreamBoothでfine-tuningしたモデルを使う前に，まず公開されている通常のStable DiffusionでDreamFusionの動作確認を行いたいと思います。

ハンバーガーの3Dモデルを生成するように学習してみます。--textのところに任意のプロンプトを入力すればよく，--workspaceには生成物などを保存するディレクトリを指定します。

$ python3 main.py --text "a hamburger" --workspace trial -O

学習が終わったら（手元の環境では90分ほどかかりました），学習済みモデルを使って360度ビデオを生成します（こちらは一瞬で終わります）。先ほど学習で用いた--workspaceと同じものを指定します。

$ python3 main.py --workspace trial -O --test

するとRGBとdepthのそれぞれでビデオが生成されます。色は毒々しいですが，形はそれらしきものができています。

RGB	Depth

DreamBoothでfine-tuningしたモデルをDreamFusionで使う

DreamFusionの動作確認が済んだところでいよいよ本題です。DreamBoothでfine-tuningすることで「いらすとや風の画像」を生成するようになった（？）Stable Diffusionのモデルを使って，DreamFusionを動かしてみます。

使用するモデルを変更するには，nerf/sd.pyにある以下の箇所を書き換えます。

class StableDiffusion(nn.Module):
    def __init__(self, device):
        super().__init__()

        # ～～中略～～
        
        # ---------------------- VAE ---------------------------------

        # 変更前
        #self.vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae", use_auth_token=self.token).to(self.device)
        
        # 変更後
        self.vae = AutoencoderKL.from_pretrained("DreamBoothで学習したモデルのパス", subfolder="vae", use_auth_token=None).to(self.device)


        # --------------------- UNet --------------------------------

        # 変更前
        #self.unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet", use_auth_token=self.token).to(self.device)
        
        # 変更後
        self.unet = UNet2DConditionModel.from_pretrained("DreamBoothで学習したモデルのパス", subfolder="unet", use_auth_token=None).to(self.device)

この状態で"sks"という識別子を含んだプロンプトを与えてDreamFusionを実行してみます。

$ python3 main.py --text "an illustration of sks girl" --workspace hoge -O

すると，以下のような3Dモデルが得られました（下の表，上段）。のっぺらぼうのようになってしまっていますし，髪型も思ってたのと違います...。

念のため，DreamBoothによるfine-tuningをしなかった際に得られる3Dモデルと比較したいと思います。プロンプトを"an illustration of girl"として生成したものを併せて表示します（下の表，下段）。こちらではミイラのような見た目のものが生成されてしまいました...。

	RGB	Depth
FTあり
FTなし

見比べてみると，fine-tuningした方は，服の色がいらすとやの画像に寄っているのかなという印象を受けます。またfine-tuningした方が輪郭などもはっきりしているように見えます。fine-tuningすることにより，ある程度似たような画像が出力されるようになったことでNeRFの学習の安定性が少し向上したのかもしれません。

プロンプトに記述することで様々なポーズを取らせることが可能なのが，DreamFusionの利点でもあります。そこで"an illustration of sks girl eating a hamburger"というプロンプトを試してみます。

$ python3 main.py --text "an illustration of sks girl eating a hamburger" --workspace fuga -O

結果は以下のようになりました。

RGB	Depth

なぜかハンバーガーが主役になってしまいました...。ハンバーガーにくっついてるピンクの塊は，ハンバーガーにしがみつく人...でしょうか？（そしてなぜかハンバーガーの色が先ほどの動作確認よりも綺麗になってる...わけわからん...。）

改良を試みる

もう少しどうにかならないものかと，以下の二つの改善策を打ってみます。

画像生成に用いるプロンプトを変えてみる
女の子以外のいらすとやの画像もDreamBoothの学習に使い汎化性能の向上を図る

まず一つ目について。今回は新たな対象物を生成できるようにしたい（例えば自分の画像を生成できるようにしたいなど）というよりは，いらすとや風の画風で生成できるようにしたいという感じなので，画像生成に使うプロンプトを以下のように少し変更してみます。

変更前：an illustration of sks girl
変更後：a sks illustration of girl

次に二つ目について。DreamBoothでは複数のインスタンスをまとめて学習することができます。そこで，いらすとやの女の子の画像だけでなく，男の子の画像と赤ちゃんの画像もまとめて学習することで汎化性能の向上を図りたいと思います。複数のインスタンスをまとめて学習するにはinstance_promptやclass_promptをまとめたconcepts_list.jsonというファイルを作成する必要があります。

[
    {
        "instance_prompt":      "a sks illustration of girl",
        "class_prompt":         "a illustration of girl",
        "instance_data_dir":    "data/irasutoya-girl",
        "class_data_dir":       "data/girl"
    },
    {
        "instance_prompt":      "a sks illustration of boy",
        "class_prompt":         "a illustration of boy",
        "instance_data_dir":    "data/irasutoya-boy",
        "class_data_dir":       "data/boy"
    },
    {
        "instance_prompt":      "a sks illustration of baby",
        "class_prompt":         "a illustration of baby",
        "instance_data_dir":    "data/irasutoya-baby",
        "class_data_dir":       "data/baby"
    }
]

そしてこれを学習の際に引数--concepts_listに渡すことで，複数インスタンスをまとめて学習することができます。

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --output_dir=$OUTPUT_DIR \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=2 --gradient_checkpointing \
  --use_8bit_adam \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=200 \
  --max_train_steps=2400 \
  --save_interval=500 \
  --concepts_list="concepts_list.json"

この方法でfine-tuningしたモデルからは以下のような画像が得られました。

prompt	girl	boy	baby
a sks illustration of ~~
a sks illustration of ~~ riding a bike
a sks illustration of ~~ eating a hamburger

赤ちゃん＞男の子＞女の子の順で上手く生成できているように見えます。それではこのモデルを使ってDreamFusionを先ほどと同じ手順で学習します。女の子，男の子，赤ちゃんのそれぞれで学習し3Dモデルを生成した結果は次のようになりました。

	girl	boy	baby
RGB
Depth

相変わらずのっぺらぼうなんですが，赤ちゃんはかなりそれっぽい3Dモデルができているように見えます（！）DreamFusionでfine-tuningした時点では赤ちゃんの画像が一番うまく生成できているように見えたので，DreamFusionも赤ちゃんの3Dモデルが一番うまくできたのは当然の結果なのかもしれません。DreamFusionによるfine-tuningがうまくいきやすいかどうかは何が関係しているのか...。

最後に，一番うまくできた赤ちゃんにサングラスをかけた3Dモデルを試しに生成してみます。プロンプトは"a sks illustration of baby wearing sunglasses"です。

RGB	Depth

四方に顔ができてしまいました...。このように複数の顔ができてしまう現象はJanus Problemと呼ばれているそうです（Janusというのは前と後ろに反対向きの二つの顔を持つローマ神話の守護神だそうです）。

We have been calling this issue where the learned 3D model has multiple faces the Janus problem (https://t.co/vauY39H9iy)
h/t @jon_barron

View-dependent prompting helps, but doesn't solve it in all cases as seen with the DreamFusion model of the squirrel below. https://t.co/Uqm9RhBVK1 pic.twitter.com/kBewdmjQNC
— Ben Poole (@poolio) October 6, 2022

おわりに

今回，DreamBoothとDreamFusionを組み合わせて，いらすとや風の3Dモデルを生成することにチャレンジしてみました。期待通りには行かない部分もありましたが，それっぽいものも得ることができました。この辺りは手法の問題なのか，私の"呪文"がヘタクソなだけなのか，現時点では定かではありませんが，改良の余地は多分にあると思われます。非常に進展の早い分野ですので，今後の進展に期待しましょう。