TRELLIS.2 Apple Silicon移植版：画像→3Dを3.5分で生成

概要

MicrosoftのTRELLIS.2画像から3Dモデルへの変換モデルを、PyTorch MPSを介してApple Silicon上でネイティブに動作するよう移植したものです。CUDAのみに依存する部分を純粋なPyTorchの代替手段に置き換えています。

主な詳細

元のTRELLIS.2は、Macでは動作しないflash_attn、nvdiffrast、およびカスタムのスパース畳み込みカーネルを必要とするCUDAを必要とします。この移植版では、それらを以下で置き換えています:

集約-分散型のスパース3D畳み込み実装 (backends/conv_none.py)
PyTorchのscaled_dot_product_attentionを使用したスパーストランスフォーマーのためのSDPAアテンション
CUDAハッシュマップ操作を置き換えるPythonベースのメッシュ抽出 (backends/mesh_extract.py)

全体の変更は9つのファイルにわたる数百行です。すべてのハードコードされた.cuda()呼び出しは、アクティブなデバイスを使用するように修正されました。

性能と要件

M4 Pro (24GB) では、単一の写真から約40万頂点のメッシュを約3.5分で生成します。生成中のメモリ使用量は、約18GBの統一メモリでピークに達します。

要件:

Apple Silicon上のmacOS (M1以降)
Python 3.11以上
24GB以上の統一メモリを推奨
モデル重み用に約15GBのディスク容量

セットアップと使用方法

クイックスタート:

git clone https://github.com/shivampkumar/trellis-mac.git
cd trellis-mac
hf auth login
bash setup.sh
source .venv/bin/activate
python generate.py path/to/image.png

HuggingFaceのゲートされたモデルへのアクセスをリクエストする必要があります: facebook/dinov3-vitl16-pretrain-lvd1689m および briaai/RMBG-2.0。

基本的な使用方法:

python generate.py photo.png
python generate.py photo.png --seed 123 --output my_model --pipeline-type 512

制限事項

テクスチャのエクスポートなし (メッシュは頂点カラーのみでエクスポートされます)
穴埋めは無効 (メッシュに小さな穴がある可能性があります)
CUDAよりも遅い (スパース畳み込みで約10倍遅い)
推論のみで、学習のサポートなし

技術的な実装

スパース3D畳み込みは、アクティブなボクセルの空間ハッシュを構築し、各カーネル位置に対して隣接特徴を集約し、行列乗算を介して重みを適用し、結果を分散加算して戻します。メッシュ抽出は、CUDAハッシュマップ操作の代わりにPython辞書を使用してflexible_dual_grid_to_meshを再実装しています。

M4 Pro (24GB)、パイプラインタイプ512でのベンチマーク: