OpenAIのShape-Eでテキストから3Dモデル生成すーる（Python、Colab）

はじめに

がちもとさんアドベントカレンダー18日目の記事です。
今日は、Shape-Eを用いてテキストから3Dモデルを生成していきます。

開発環境

• Windows 11 PC
• Python 3.11

導入

1．Colabからコピー

2．Googleドライブに接続

from google.colab import drive
drive.mount('/content/drive')

3．GitHubからShape-Eをクローン

!git clone https://github.com/openai/shap-e

4．ライブラリのインストール

%cd /content/shap-e/
!pip install -e .

5．メインのプログラムを作成

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, gif_widget

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

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

batch_size = 4
guidance_scale = 15.0
prompt = "a pink shark" #@param {type:"string"}

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 = 128 #@param {type:"number"}

cameras = create_pan_cameras(size, device)
for i, latent in enumerate(latents):
    images = decode_latent_images(xm, latent, cameras, rendering_mode=render_mode)
    display(gif_widget(images))

# Example of saving the latents as meshes.
from shap_e.util.notebooks import decode_latent_mesh

for i, latent in enumerate(latents):
    with open(f'example_mesh_{i}.ply', 'wb') as f:
        decode_latent_mesh(xm, latent).tri_mesh().write_ply(f)

実行結果

prompt: "a pink shark"
size: 128

生成した3Dモデル（PLY）をダウンロードし、MeshLabで表示

お疲れさまでした。