[DreamFusion] Stable Diffusionを用いてテキストから3Dモデルを生成する

2022年10月12日水曜日

Artificial Intelligence

本記事では、DreamFusionと呼ばれる機械学習手法を用いて、テキストから3Dモデルを生成する方法をご紹介します。

アイキャッチ

DreamFusion

概要

DreamFusionは、事前トレーニング済みのText to Imageタスクの拡散モデルを使用して、Text to 3D modelタスクを実現する3次元モデル合成技術です。

2022年現在、テキストから画像を生成するText to Imageタスクにおいて、大きなブレークスルーを実現している拡散モデルは、何十億もの画像とテキストのペアでトレーニングされています。
この手法をText to 3D modelタスクに適用する場合、大規模なラベル付けされた3Dデータが必要ですが、これらの用意は現状困難です。

このためDreamFusionでは、Deep Dreamのような手順でランダムな角度からのレンダリングの損失が最小化するように勾配降下によって3Dモデルを最適化しています。
この方法により、3Dデータを必要とせず3次元モデルの生成を可能にしています。

アーキテクチャ

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

本記事では上記手法を用いて、テキストから3Dモデルを生成する方法をご紹介します。

デモ(Colaboratory)

それでは、実際に動かしながらテキストから3Dモデルを生成していきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

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

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

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更

初めにGithubからソースコードを取得します。

%cd /content

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

# Commits on Oct 10, 2022
%cd /content/stable-dreamfusion
!git checkout 76ff91ebc60f99249c0f828cf70cebba369e2cd4

次にライブラリをインストールします。

%cd /content/stable-dreamfusion

# install requirements
! pip install -r requirements.txt
! pip install git+https://github.com/NVlabs/nvdiffrast/@fad71a4ad3dc8530a4eeaeb55a65b2e8c127d69d

# install extension modules
! bash scripts/install_ext.sh

次にHuggingFaceからStable Diffusionの事前学習済みモデルをロードするためHugginFaceにログインします。
AccessTokenの取得方法は以下の記事をご参照ください。

from huggingface_hub import notebook_login
from google.colab import output

output.enable_custom_widget_manager()
notebook_login()

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

トレーニングのセットアップ

任意のPromptに応じた3Dモデルを生成するためのトレーニングをセットアップしていきます。
今回はbananaの3Dモデルを生成してみたいと思います。

Prompt_text = "a DSLR photo of a banana, low ploy" #@param {type: 'string'}

Training_iters = 5000 #@param {type: 'integer'}
Learning_rate = 1e-3 #@param {type: 'number'}
Training_nerf_resolution = 64  #@param {type: 'integer'}
Seed = 12 #@param {type: 'integer'}
Lambda_entropy = 1e-4 #@param {type: 'number'}
Checkpoint = 'scratch' #@param {type: 'string'}

# 出力先
Workspace = "trial" #@param{type: 'string'}
# processings
Prompt_text = "'" + Prompt_text + "'"

以下の通りトレーニングを実行します。

GPUにもよりますが、1エポック1分ほどなので、軽く見積もっても1時間程度を要します。
本記事では、時間短縮のため20エポックで切り上げています。

%cd /content/stable-dreamfusion

%run main.py \
  -O \
  --text {Prompt_text} \
  --workspace {Workspace} \
  --iters {Training_iters} \
  --lr {Learning_rate} \
  --w {Training_nerf_resolution} \
  --h {Training_nerf_resolution} \
  --seed {Seed} \
  --lambda_entropy {Lambda_entropy} \
  --ckpt {Checkpoint} \
  --save_mesh

Text to 3D model

学習途中のモデルではありますが、3Dモデルを生成してみます。

%cd /content/stable-dreamfusion

%run main.py \
  -O \
  --test \
  --workspace {Workspace} \
  --save_mesh

20エポックでの出力結果は以下の通りです。
残念ながらこの段階ではまだバナナの視認できるような形状ではありません。
その他プロンプトによっては50エポックでもうまく3Dモデルが生成されない場合もあり、やや不安定です。

うまく生成され次第プロンプトと生成結果を更新予定です。

import os
import glob
from IPython.display import HTML
from base64 import b64encode

def get_latest_file(path):
  dir_list = glob.glob(path)
  dir_list.sort(key=lambda x: os.path.getmtime(x))
  return dir_list[-1]

def show_video(video_path, video_width = 600):
   
  video_file = open(video_path, "r+b").read()
  video_url = f"data:video/mp4;base64,{b64encode(video_file).decode()}"

  return HTML(f"""<video width={video_width} controls><source src="{video_url}"></video>""")
 
rgb_video = get_latest_file(os.path.join(Workspace, 'results', '*_rgb.mp4'))
show_video(rgb_video)
result

まとめ

本記事では、DreamFusionを用いてText to 3D modelを試してみました。
現状、Google Colaboratoryでは学習コストにやや難があります。トレーニングの高速化などを探りたいところです。

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


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

参考文献

1.  論文 - DreamFusion: Text-to-3D using 2D Diffusion

2. GitHub - ashawkey/stable-dreamfusion

AIで副業ならココから!

まずは無料会員登録

プロフィール

メーカーで研究開発を行う現役エンジニア
組み込み機器開発や機会学習モデル開発に従事しています

本ブログでは最新AI技術を中心にソースコード付きでご紹介します


Twitter

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology