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

# はじめに

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

昨日(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](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3506235/e6207c88-766d-59e3-45b8-4ed51e5e96bb.png)

公式の[Github](https://github.com/NVIDIA/Cosmos-Tokenizer)
公式の[HuggingFace](https://huggingface.co/nvidia/Cosmos-1.0-Tokenizer-CV8x8x8)

# 手順
## ライブラリのインストール

```bash
# 仮想環境の作成と有効化
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`スクリプトを作成して以下の内容を記述：

```python
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)
```

以下のコマンドで実行：

```bash
python3 download.py
```

---

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

```python
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`)

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

```python
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}")
```

実行結果：
```bash
Latent Shape: torch.Size([1, 16, 2, 64, 64])
Reconstructed Shape: torch.Size([1, 3, 9, 512, 512])
```

---

##### 2. Discrete Tokenizer (`test02.py`)

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

```python
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}")
```

実行結果：
```bash
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])
```

---

#### 動画ファイルのテスト

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

```bash
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
```

処理結果：
```bash
Found 1 videos from test_data/*.mp4.
Outputting test_data/reconstructions/video.mp4 ...
```

PSNR（画質評価）の計算：

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

結果：
```bash
(.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で画質評価**  
   動画の復元精度を簡単に評価できる。

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