Shap-Eを使ってテキストから3Dオブジェクトを生成する

github.com

はじめに

前回「Stable-Dreamfusion」というのを使って3D生成を行いました。
touch-sp.hatenablog.com
今回は「Shap-E」というのを使わせて頂きました。

環境

Ubuntu 22.04 on WSL2
CUDA 11.8
Python 3.10

導入

pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --index-url https://download.pytorch.org/whl/cu118
pip install pyyaml
pip install ipywidgets
pip install git+https://github.com/facebookresearch/pytorch3d.git
git clone https://github.com/openai/shap-e
cd shap-e
pip install -e .

上記インストールで失敗する場合には最初に以下を実行すれば解決すると思います。

python -m pip install --upgrade pip
pip install -U setuptools wheel

Pythonスクリプト

import torch

from shap_e.diffusion.sample import sample_latents
from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
from shap_e.models.download import load_model, load_config
from shap_e.util.notebooks import create_pan_cameras, decode_latent_images

device = torch.device('cuda')

xm = load_model('transmitter', device=device)
model = load_model('text300M', device=device)
diffusion = diffusion_from_config(load_config('diffusion'))

batch_size = 2
guidance_scale = 15.0
prompt = "A penguin"

latents = sample_latents(
    batch_size=batch_size,
    model=model,
    diffusion=diffusion,
    guidance_scale=guidance_scale,
    model_kwargs=dict(texts=[prompt] * batch_size),
    progress=True,
    clip_denoised=True,
    use_fp16=True,
    use_karras=True,
    karras_steps=64,
    sigma_min=1e-3,
    sigma_max=160,
    s_churn=0,
)

render_mode = 'nerf'    # you can change this to 'stf'
size = 256              # this is the size of the renders; higher values take longer to render.

cameras = create_pan_cameras(size, device)
for i, latent in enumerate(latents):
    images = decode_latent_images(xm, latent, cameras, rendering_mode=render_mode)
    images[0].save(
        f"result_{i}.gif", format="GIF", save_all=True, append_images=images[1:], duration=100, loop=0)

結果



3Dオブジェクトファイル(拡張子obj)の作り方

先ほどのPythonスクリプトの後半を以下に書き換えれば可能です。

from shap_e.util.notebooks import decode_latent_mesh

for i, latent in enumerate(latents):
    t = decode_latent_mesh(xm, latent).tri_mesh()
    with open(f'example_mesh_{i}.obj', 'w') as f:
        t.write_obj(f)



以下はBlenderで読み込んだところです。

上のように表示させるには「オブジェクトモード」を「頂点ペイント」に変更する必要があります。

VRAM使用量は

VRAM 12GBのRTX 3080で実行可能でした。

さいごに

今回はテキストからの3Dオブジェクト生成でした。

画像からの3Dオブジェクト生成はこちらです。
touch-sp.hatenablog.com