【Diffusers】Diffusers から OpenAI の Shap-E が使えるようになったのでさっそくテキストから3Dオブジェクトを生成してみました

最終更新日：2023年7月21日

はじめに

「Shap-E」は以前にこちらの記事で扱いました。
touch-sp.hatenablog.com
Diffuses 0.18.0からShap-Eがサポートされたようなのでさっそく使ってみました。

Pythonスクリプト

たったこれだけです。ずいぶん使いやすくなっています。

モデルは自動的にダウンロードされます。事前準備は不要です。

from diffusers import DiffusionPipeline
from diffusers.utils import export_to_gif

ckpt_id = "openai/shap-e"
pipe = DiffusionPipeline.from_pretrained(ckpt_id).to("cuda")

images = pipe(
    prompt="A penguin",
    guidance_scale=15.0,
    num_images_per_prompt=1,
    num_inference_steps=64,
    frame_size=512,
).images

for i, image in enumerate(images):
    export_to_gif(image, f"penguin_{i}.gif")

動作環境

Windows 11
CUDA 11.7
Python 3.10

インストールするものも最低限で動作します。

pip install torch==2.0.1+cu117 --index-url https://download.pytorch.org/whl/cu117
pip install diffusers[torch]==0.18.2
pip install transformers==4.31.0

VRAM使用量

float32

上記スクリプトを実行すると9.1GBのVRAMを使用していました。

float16

float16に変更してみました。

import torch
from diffusers import DiffusionPipeline
from diffusers.utils import export_to_gif

ckpt_id = "openai/shap-e"
pipe = DiffusionPipeline.from_pretrained(ckpt_id, torch_dtype=torch.float16).to("cuda")

images = pipe(
    prompt="A penguin",
    guidance_scale=15.0,
    num_images_per_prompt=1,
    num_inference_steps=64,
    frame_size=512,
).images

for i, image in enumerate(images):
    export_to_gif(image, f"penguin_{i}.gif")

5.3GBまで使用量を減らすことが出来ました。