WSL2でOpen OASISを試してみる

2024年11月2日 17:34

「Etched社とDecart社の共同開発により、AIがプレイヤーの入力に応じて動画を生成する初のインタラクティブなオープンワールドゲームモデル」、「ユーザーのキーボードやマウスの操作に基づき、物理演算やゲームルール、グラフィックスを内部でシミュレートし、フレームごとにゲーム映像を生成」、「従来のゲームエンジンに代わるAI駆動の未来を目指す研究の第一歩と位置付けられている」らしいOASISのOpenモデルが公開されたので試してみます。

将来的には、Etched社のTransformer ASIC「Sohu」チップなるものの上で4Kで動くぞー、ということらしいので、まー、資金（ry。

使用するPCはドスパラさんの「GALLERIA UL9C-R49」。スペックは
・CPU: Intel® Core™ i9-13900HX Processor
・Mem: 64 GB
・GPU: NVIDIA® GeForce RTX™ 4090 Laptop GPU(16GB)
・GPU: NVIDIA® GeForce RTX™ 4090 (24GB)
・OS: Ubuntu22.04 on WSL2（Windows 11）
です。

1. 準備

環境構築

python3 -m venv open-oasis
cd $_
source bin/activate

リポジトリをクローンして、インストールします。

git clone https://github.com/etched-ai/open-oasis
cd open-oasis

パッケージのインストールです。

pip install -r requirements.txt

モデルのダウンロード

モデルをダウンロードして、カレントディレクトリにリンクを作っておきます。

huggingface-cli download Etched/oasis-500m oasis500m.pt # DiT checkpoint
huggingface-cli download Etched/oasis-500m vit-l-20.pt  # ViT VAE checkpoint
#
ln -s ~/.cache/huggingface/hub/models--Etched--oasis-500m/snapshots/e562e3ac8e99c11b9329ad0d057cb2cd69fee20f/oasis500m.pt .
ln -s ~/.cache/huggingface/hub/models--Etched--oasis-500m/snapshots/e562e3ac8e99c11b9329ad0d057cb2cd69fee20f/vit-l-20.pt .

2. 試してみる

generate.pyがサンプルコードとして提供されています。
einops使ってrearrageeで書いていただくとテンソル演算で何しているのかわかりやすくて良いですね。

device = "cuda:0" とハードコードされていましたが、手が勝手に動くのでCUDA_VISIBLE_DEVICESも合わせて実行です。

CUDA_VISIBLE_DEVICES=0 python generate.py

VRAMの消費量は、6.0GB付近。

待つこと約3分20秒でカレントディレクトリに video.mp4 が生成されました。32フレームです。

100%|███████████████████████████████████████████████████████████████████████████████████████████| 31/31 [03:19<00:00,  6.44s/it]
generation saved to video.mp4.

うん…？　リアル・・・タイム？？？

open-oasisを試している。https://t.co/cyzkhNV9px
4090(24GB)で32フレーム生成に3分22秒ほど。VRAMは5.9GBほど。 pic.twitter.com/oqTGNDZeHz
— NOGUCHI, Shoji (@noguchis) November 1, 2024

generate.pyの38行目の値を

$ grep -n ^video_id generate.py 
38:video_id = "snippy-chartreuse-mastiff-f79998db196d-20220401-224517.chunk_001"
$

sample_dataディレクトにある以下のファイルのベースネームに置き換えれば、その.mp4を入力にしてvideo.mpが生成できます。

$ ls sample_data/*.mp4
sample_data/Player729-f153ac423f61-20210806-224813.chunk_000.mp4
sample_data/snippy-chartreuse-mastiff-f79998db196d-20220401-224517.chunk_001.mp4
sample_data/treechop-f153ac423f61-20210916-183423.chunk_000.mp4
$

それが、こちら。

open-oasis続き。
サンプル動画が3つあったので、https://t.co/MDLTu598eVをごにょごにょして残りの2つも生成。 pic.twitter.com/VRUlKsJv66
— NOGUCHI, Shoji (@noguchis) November 2, 2024

4. まとめ

リソース使用状況を中心にまとめておきます。

32フレーム生成時、GPU VRAM使用量は6GBほど。
生成時間は3分20秒ほど。

まー、現時点では技術的なデモンストレーションですね。なんだか、Soraを思い出してしまった。