LoginSignup
5
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@syun88

CES2025 NVIDIA Cosmosと関係ないモデルをCosmos-Tokenizer: 高性能なトークン化と復元技術

Last updated at Posted at 2025-01-07

はじめに

こんにちは、しゅんです！

昨日(2025/01/07日本時間11:30)のCES2025ライブで、NVIDIAのCEO Jensen HuangがCosmosをオープンソース化すると発表しました。それを聞いて試してみたくなった。

お詫び、このモデルは去年年末で更新したモデル、正確に言うと今回のCES2025とあまり関係ないです。

この記事では、公式のdocker環境ではなくローカルで構築しています。Cosmos-Tokenizerのインストールから動作確認までの手順をまとめています。

Cosmos-Tokenizerとは？

Cosmos-Tokenizerは、画像や動画を効率的にトークン化（符号化）し、復元するためのNVIDIAの最新モデルです。このモデルは連続的（Continuous）および離散的（Discrete）なトークン化をサポートしており、機械学習の前処理や圧縮に役立ちます。GPTに聞いてみたが、多分そういうこと、正直言ってまだ理解できていない。

  • データ圧縮: 高精度でデータサイズを圧縮。
  • 機械学習の前処理: モデルの入力データを最適化。
  • 生成AI: トークンを用いた画像や動画生成。

Cosmos-Tokenizerの特徴

  1. 連続（Continuous）と離散（Discrete）モードのサポート

    • Continuous Tokenizer: 入力データを潜在空間に変換し、数値ベースで圧縮。
    • Discrete Tokenizer: トークンとして符号化し、情報量を極限まで削減。

  2. マルチモーダル対応
    画像（2D）や動画（3D）の処理が可能で、トークン化と復元を高速に行います。

  3. 効率的なアーキテクチャ

    • NVIDIAの最新GPUに最適化。
    • CUDAを活用した高速な処理。

  4. 柔軟なトークンサイズ
    トークンのサイズや次元数を選択可能（例: 8x8x8, 16x16x16など）。

  5. オープンソースと事前学習済みモデル
    Hugging Faceリポジトリから事前学習済みのモデルをダウンロード可能。

Cosmos-Tokenizerの仕組み

Cosmos-Tokenizerは以下のような流れで動作します：

  1. 符号化（エンコード）
    入力データをトークンまたは潜在ベクトルに変換します。これによりデータサイズが削減されます。

    • Continuous: 潜在空間の表現を生成。
    • Discrete: 離散トークンを生成。

  2. 復元（デコード）
    トークンまたは潜在ベクトルを元の形式に再構成します。

  3. トークンの利用

    • 圧縮データの伝送や保存。
    • 機械学習モデルへの入力。
    • 復元による評価（PSNRやSSIMなど）。

主な利用例

  1. データ圧縮
    高精度を維持しながら、動画や画像のファイルサイズを削減。特に帯域幅や保存容量が限られた環境で有効。

  2. 生成AI
    トークンを使用して新しい画像や動画を生成。

  3. 機械学習の効率化
    トークン化されたデータを直接モデルに入力することで、学習と推論の速度が向上。

Cosmos-Tokenizerの主なバージョン

Cosmos-Tokenizerは複数のバージョンを提供しており、それぞれ異なるデータタイプや用途に特化しています。

モデル名 特徴 用途
Cosmos-0.1-Tokenizer-CV8x8x8 Continuousモデル (動画対応) 高精度な動画圧縮・復元
Cosmos-0.1-Tokenizer-DV8x8x8 Discreteモデル (動画対応) トークンベースの動画処理
Cosmos-0.1-Tokenizer-CI8x8 Continuousモデル (画像対応) 高精度な画像圧縮・復元
Cosmos-0.1-Tokenizer-DI8x8 Discreteモデル (画像対応) トークンベースの画像処理
Cosmos-1.0-Tokenizer-CV8x8x8 Continuousモデル (改良版) 高速な動画処理

image.png

公式のGithub
公式のHuggingFace

手順

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

# 仮想環境の作成と有効化
python3 -m venv .venv
source .venv/bin/activate

# Cosmos-Tokenizerのリポジトリをクローン
git clone https://github.com/NVIDIA/Cosmos-Tokenizer.git
cd Cosmos-Tokenizer

# 依存関係をインストール
pip3 install -r requirements.txt

# 必要なパッケージを追加でインストール
sudo apt install -y ffmpeg
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install --upgrade mediapy

モデルのダウンロード

NVIDIAの提供するHugging Faceリポジトリからモデルをダウンロードします。

download.pyスクリプトを作成して以下の内容を記述：

from huggingface_hub import login, snapshot_download
import os

# Hugging Faceトークンを設定
login(token="YOUR_HUGGING_FACE_TOKEN", add_to_git_credential=True)

# モデルリスト
model_names = [
    "Cosmos-0.1-Tokenizer-CV8x8x8",
    "Cosmos-0.1-Tokenizer-DV8x8x8",
    # 必要に応じて追加
]

# モデルをローカルに保存
for model_name in model_names:
    hf_repo = "nvidia/" + model_name
    local_dir = "pretrained_ckpts/" + model_name
    os.makedirs(local_dir, exist_ok=True)
    print(f"Downloading {model_name}...")
    snapshot_download(repo_id=hf_repo, local_dir=local_dir)

以下のコマンドで実行：

python3 download.py

公式のサンプルコードを少し修正した。

import torch
from cosmos_tokenizer.video_lib import CausalVideoTokenizer
# model_name = "Cosmos-Tokenizer-1.0-CV8x8x8"

model_name = "Cosmos-0.1-Tokenizer-CV8x8x8"
input_tensor = torch.randn(1, 3, 9, 512, 512).to('cuda').to(torch.bfloat16)  # [B, C, T, H, W]
encoder = CausalVideoTokenizer(checkpoint_enc=f'pretrained_ckpts/{model_name}/encoder.jit')
(latent,) = encoder.encode(input_tensor)
torch.testing.assert_close(latent.shape, (1, 16, 2, 64, 64))

# The input tensor can be reconstructed by the decoder as:
decoder = CausalVideoTokenizer(checkpoint_dec=f'pretrained_ckpts/{model_name}/decoder.jit')
reconstructed_tensor = decoder.decode(latent)
torch.testing.assert_close(reconstructed_tensor.shape, input_tensor.shape)

GPTが作ったコードの実行

まじでまだ理解している段階で今回はただ試してただけ

1. Continuous Tokenizer (test01.py)

連続的なトークン化と復元を実行するコード：

import torch
from cosmos_tokenizer.video_lib import CausalVideoTokenizer

model_name = "Cosmos-0.1-Tokenizer-CV8x8x8"
encoder_path = f"pretrained_ckpts/{model_name}/encoder.jit"
decoder_path = f"pretrained_ckpts/{model_name}/decoder.jit"

encoder = CausalVideoTokenizer(checkpoint_enc=encoder_path)
decoder = CausalVideoTokenizer(checkpoint_dec=decoder_path)

input_tensor = torch.randn(1, 3, 9, 512, 512).to('cuda').to(torch.bfloat16)

(latent,) = encoder.encode(input_tensor)
print(f"Latent Shape: {latent.shape}")

reconstructed_tensor = decoder.decode(latent)
print(f"Reconstructed Shape: {reconstructed_tensor.shape}")

実行結果：

Latent Shape: torch.Size([1, 16, 2, 64, 64])
Reconstructed Shape: torch.Size([1, 3, 9, 512, 512])
2. Discrete Tokenizer (test02.py)

離散的なトークン化と復元を実行するコード：

import torch
from cosmos_tokenizer.video_lib import CausalVideoTokenizer

model_name = "Cosmos-0.1-Tokenizer-DV8x8x8"
encoder_path = f"pretrained_ckpts/{model_name}/encoder.jit"
decoder_path = f"pretrained_ckpts/{model_name}/decoder.jit"

encoder = CausalVideoTokenizer(checkpoint_enc=encoder_path)
decoder = CausalVideoTokenizer(checkpoint_dec=decoder_path)

input_tensor = torch.randn(1, 3, 9, 512, 512).to('cuda').to(torch.bfloat16)

(indices, codes) = encoder.encode(input_tensor)
print(f"Indices Shape: {indices.shape}")
print(f"Codes Shape: {codes.shape}")

reconstructed_tensor = decoder.decode(indices)
print(f"Reconstructed Shape: {reconstructed_tensor.shape}")

実行結果：

Indices Shape: torch.Size([1, 2, 64, 64])
Codes Shape: torch.Size([1, 6, 2, 64, 64])
Reconstructed Shape: torch.Size([1, 3, 9, 512, 512])

動画ファイルのテスト

以下のコマンドで動画データを処理：

python3 -m cosmos_tokenizer.video_cli \
    --video_pattern 'test_data/*.mp4' \
    --checkpoint_enc pretrained_ckpts/Cosmos-0.1-Tokenizer-CV8x8x8/encoder.jit \
    --checkpoint_dec pretrained_ckpts/Cosmos-0.1-Tokenizer-CV8x8x8/decoder.jit

処理結果：

Found 1 videos from test_data/*.mp4.
Outputting test_data/reconstructions/video.mp4 ...

PSNR（画質評価）の計算：

ffmpeg -i test_data/video.mp4 -i test_data/reconstructions/video.mp4 -lavfi psnr="stats_file=psnr.log" -f null -

結果：

(.venv) syun@syun:/media/syun/ssd02/python_learning/nvidia/cosmos/Cosmos-Tokenizer$ ffmpeg -i test_data/video.mp4 -i test_data/reconstructions/video.mp4 -lavfi psnr="stats_file=psnr.log" -f null -
ffmpeg version 6.1.1-3ubuntu5+esm2 Copyright (c) 2000-2023 the FFmpeg developers
  built with gcc 13 (Ubuntu 13.2.0-23ubuntu4)
  configuration: --prefix=/usr --extra-version=3ubuntu5+esm2 --toolchain=hardened --libdir=/usr/lib/x86_64-linux-gnu --incdir=/usr/include/x86_64-linux-gnu --arch=amd64 --enable-gpl --disable-stripping --disable-omx --enable-gnutls --enable-libaom --enable-libass --enable-libbs2b --enable-libcaca --enable-libcdio --enable-libcodec2 --enable-libdav1d --enable-libflite --enable-libfontconfig --enable-libfreetype --enable-libfribidi --enable-libglslang --enable-libgme --enable-libgsm --enable-libharfbuzz --enable-libmp3lame --enable-libmysofa --enable-libopenjpeg --enable-libopenmpt --enable-libopus --enable-librubberband --enable-libshine --enable-libsnappy --enable-libsoxr --enable-libspeex --enable-libtheora --enable-libtwolame --enable-libvidstab --enable-libvorbis --enable-libvpx --enable-libwebp --enable-libx265 --enable-libxml2 --enable-libxvid --enable-libzimg --enable-openal --enable-opencl --enable-opengl --disable-sndio --enable-libvpl --disable-libmfx --enable-libdc1394 --enable-libdrm --enable-libiec61883 --enable-chromaprint --enable-frei0r --enable-ladspa --enable-libbluray --enable-libjack --enable-libpulse --enable-librabbitmq --enable-librist --enable-libsrt --enable-libssh --enable-libsvtav1 --enable-libx264 --enable-libzmq --enable-libzvbi --enable-lv2 --enable-sdl2 --enable-libplacebo --enable-librav1e --enable-pocketsphinx --enable-librsvg --enable-libjxl --enable-shared
  libavutil      58. 29.100 / 58. 29.100
  libavcodec     60. 31.102 / 60. 31.102
  libavformat    60. 16.100 / 60. 16.100
  libavdevice    60.  3.100 / 60.  3.100
  libavfilter     9. 12.100 /  9. 12.100
  libswscale      7.  5.100 /  7.  5.100
  libswresample   4. 12.100 /  4. 12.100
  libpostproc    57.  3.100 / 57.  3.100
Input #0, mov,mp4,m4a,3gp,3g2,mj2, from 'test_data/video.mp4':
  Metadata:
    major_brand     : qt  
    minor_version   : 0
    compatible_brands: qt  
    creation_time   : 2024-11-06T06:51:35.000000Z
    com.apple.quicktime.location.accuracy.horizontal: 5.780658
    com.apple.quicktime.full-frame-rate-playback-intent: 0
    com.apple.quicktime.make: Apple
    com.apple.quicktime.model: iPhone 15 Pro Max
    com.apple.quicktime.software: 18.1
    com.apple.quicktime.creationdate: 2024-11-02T22:20:00-0500
    com.apple.photos.originating.signature: AZcvUsJG1u14NTSnwxpJi4UaLJTo
  Duration: 00:00:02.00, start: 0.000000, bitrate: 14316 kb/s
  Stream #0:0[0x1](und): Video: h264 (High) (avc1 / 0x31637661), yuvj420p(pc, bt709, progressive), 882x1568, 14296 kb/s, 60.15 fps, 59.94 tbr, 600 tbn (default)
    Metadata:
      creation_time   : 2024-11-06T06:51:35.000000Z
      handler_name    : Core Media Video
      vendor_id       : [0][0][0][0]
      encoder         : H.264
  Stream #0:1[0x2](und): Data: none (mebx / 0x7862656D), 0 kb/s (default)
    Metadata:
      creation_time   : 2024-11-06T06:51:35.000000Z
      handler_name    : Core Media Metadata
  Stream #0:2[0x3](und): Data: none (mebx / 0x7862656D), 0 kb/s (default)
    Metadata:
      creation_time   : 2024-11-06T06:51:35.000000Z
      handler_name    : Core Media Metadata
Input #1, mov,mp4,m4a,3gp,3g2,mj2, from 'test_data/reconstructions/video.mp4':
  Metadata:
    major_brand     : isom
    minor_version   : 512
    compatible_brands: isomiso2avc1mp41
    encoder         : Lavf60.16.100
  Duration: 00:00:05.00, start: 0.000000, bitrate: 1356 kb/s
  Stream #1:0[0x1](und): Video: h264 (High) (avc1 / 0x31637661), yuv420p(progressive), 882x1568, 1352 kb/s, 24 fps, 24 tbr, 12288 tbn (default)
    Metadata:
      handler_name    : VideoHandler
      vendor_id       : [0][0][0][0]
      encoder         : Lavc60.31.102 libx264
Stream mapping:
  Stream #0:0 (h264) -> psnr
  Stream #1:0 (h264) -> psnr
  psnr:default -> Stream #0:0 (wrapped_avframe)
Press [q] to stop, [?] for help
[swscaler @ 0x6090d6ffcf80] deprecated pixel format used, make sure you did set range correctly
[Parsed_psnr_0 @ 0x6090d55d48c0] not matching timebases found between first input: 1/600 and second input 1/12288, results may be incorrect!
Output #0, null, to 'pipe:':
  Metadata:
    major_brand     : qt  
    minor_version   : 0
    compatible_brands: qt  
    com.apple.photos.originating.signature: AZcvUsJG1u14NTSnwxpJi4UaLJTo
    com.apple.quicktime.location.accuracy.horizontal: 5.780658
    com.apple.quicktime.full-frame-rate-playback-intent: 0
    com.apple.quicktime.make: Apple
    com.apple.quicktime.model: iPhone 15 Pro Max
    com.apple.quicktime.software: 18.1
    com.apple.quicktime.creationdate: 2024-11-02T22:20:00-0500
    encoder         : Lavf60.16.100
  Stream #0:0: Video: wrapped_avframe, yuvj420p(pc, bt709, progressive), 882x1568, q=2-31, 200 kb/s, 59.94 fps, 59.94 tbn
    Metadata:
      encoder         : Lavc60.31.102 wrapped_avframe
frame=    0 fps=0.0 q=-0.0 size=       0kB time=00:00:00.00 bitrate=N/A spe[out#0/null @ 0x6090d6ffb700] video:90kB audio:0kB subtitle:0kB other streams:0kB global headers:0kB muxing overhead: unknown
frame=  192 fps=0.0 q=-0.0 Lsize=N/A time=00:00:04.95 bitrate=N/A speed=15.1x    
[Parsed_psnr_0 @ 0x6090d55d48c0] PSNR y:19.142809 u:33.280030 v:40.252139 average:20.853721 min:18.597245 max:44.119122

Cosmos-Tokenizerの将来性
NVIDIAのCosmos-Tokenizerは、データ圧縮や生成AIの分野で非常に有望な技術です。オープンソース化により、多くの開発者がこの技術を活用し、さらなる発展が期待されています。特に、動画処理やAIアプリケーションの効率化において、圧倒的なパフォーマンスを発揮します。 これはGPTが
PSNR y:19.141534 u:33.283610 v:40.250060 average:20.852491この結果を見てこの結論を出しました。

まとめ

今回、Cosmos-Tokenizerを試した結果、以下のことが分かりました：

  1. 動作確認が可能
    Continuous/Discrete両方のモデルでトークン化と復元が問題なく行えることを確認。

  2. PSNRで画質評価
    動画の復元精度を簡単に評価できる。

この記事が参考になれば嬉しいです！最後までお読みいただきありがとうございました。

5
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
5
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?