NVIDIA Cosmos Tokenizer：视频与图像的神经网络分词器

一款让视频"瘦身"2048倍的黑科技，NVIDIA 开源视觉 tokenization 利器

想象一下，你要训练一个能生成 10 秒高画质视频的 AI 模型。以 1080p、30fps 计算，这短短 10 秒就包含 9000 帧画面、接近 6GB 的原始像素数据——直接塞进 Transformer 模型根本不现实。业界通用的解决方案是先将视频"压缩"为少量的"关键词"（token），再用生成模型处理这些关键词。NVIDIA Cosmos Tokenizer 就是这个压缩环节的顶尖选手：它能将视频压缩 2048 倍，同时保持画质基本不变，且速度比同类方案快 12 倍。

2025 年 1 月，NVIDIA 在发布 Cosmos 世界模型平台时同步开源了 Cosmos Tokenizer，很快便在 GitHub 收获超过 1700 颗星，成为视频生成管线中最受欢迎的预处理工具之一。同年 2 月 11 日，该仓库转为只读归档，后续开发已迁移至 NVIDIA/Cosmos 主仓库。

图：Cosmos Tokenizer 整体架构，编码器将视频帧压缩为潜在表示，解码器完成重建

背景：为什么视频 tokenization 至关重要

2023—2024 年间，Runway Gen-2/Sora 等视频生成模型相继爆火，背后有一个共同的技术基础：视频压缩 tokenization。在 Stable Diffusion 时代，图像可以用 VAE 压缩 8 倍；到了视频时代，单帧压缩已经不够用了——必须同时在空间（图像分辨率）和时间（帧间冗余）两个维度做压缩，才能让 Transformer 处理得起。

Cosmos Tokenizer 正是在这一背景下诞生。NVIDIA 团队从 Cosmos 物理 AI 项目中提炼出这套通用视觉分词能力，目标是成为所有视频生成模型的标准前端。无论你是做文生视频、机器人操控视频仿真，还是自动驾驶数据压缩，Cosmos Tokenizer 都能提供一个开箱即用的高质量压缩层。

从技术谱系上看，Cosmos Tokenizer 继承并改进了 MagViT（Google）和 SOTA-VAE 系列方法，核心引入了因果卷积（causal convolution）来处理时序依赖，并通过 FSQ（Finite Scalar Quantization） 替代传统的 VQ（Vector Quantization）来减少量化误差。

核心功能：12 款预训练模型，覆盖连续与离散两条路线

Cosmos Tokenizer 提供 12 款预训练模型，分别从两个维度切分：

类型	图像版（I）	视频版（V）
连续版（C）	CI8x8、CI16x16	CV4x8x8、CV8x8x8、CV8x16x16
离散版（D）	DI8x8、DI16x16	DV4x8x8、DV8x8x8、DV8x16x16

命名规则清晰易懂：数字代表压缩倍率，CV4x8x8 表示时间压缩 4 倍、空间压缩 8×8。连续版（CV/DI 等）输出浮点 latent tensor，离散版（DV/DI 等）输出整数 token ID，后者更适合搭配离散自回归模型（如 LLM 风格的语言模型做视频生成）。

使用示例（连续视频分词）：

import torch
from cosmos_tokenizer.video_lib import CausalVideoTokenizer

model_name = "Cosmos-0.1-Tokenizer-CV4x8x8"
encoder = CausalVideoTokenizer(
    checkpoint_enc=f'pretrained_ckpts/{model_name}/encoder.jit')
decoder = CausalVideoTokenizer(
    checkpoint_dec=f'pretrained_ckpts/{model_name}/decoder.jit')

# 输入: [B=1, C=3, T=9帧, H=512, W=512]
input_tensor = torch.randn(1, 3, 9, 512, 512).to('cuda').to(torch.bfloat16)
(latent,) = encoder.encode(input_tensor)
# latent shape: [1, 16, 3, 64, 64] — 压缩了 64 倍（空间 8x8）
reconstructed = decoder.decode(latent)
torch.testing.assert_close(reconstructed.shape, input_tensor.shape)

使用示例（离散视频分词）：

model_name = "Cosmos-0.1-Tokenizer-DV4x8x8"
encoder = CausalVideoTokenizer(
    checkpoint_enc=f'pretrained_ckpts/{model_name}/encoder.jit')
decoder = CausalVideoTokenizer(
    checkpoint_dec=f'pretrained_ckpts/{model_name}/decoder.jit')

(indices, codes) = encoder.encode(input_tensor)
# indices: [1, 3, 64, 64] — 每个空间位置一个离散的 token ID（范围 1~64K）
# codes: [1, 6, 3, 64, 64] — 量化前的连续表示
reconstructed = decoder.decode(indices)

此外，Cosmos Tokenizer 还提供了 CLI 工具，一行命令完成批量视频/图像分词：

# 图像分词
python3 -m cosmos_tokenizer.image_cli \
    --image_pattern 'test_data/*.png' \
    --checkpoint_enc pretrained_ckpts/Cosmos-0.1-Tokenizer-CI8x8/encoder.jit \
    --checkpoint_dec pretrained_ckpts/Cosmos-0.1-Tokenizer-CI8x8/decoder.jit

# 视频分词
python3 -m cosmos_tokenizer.video_cli \
    --video_pattern 'test_data/video.mp4' \
    --checkpoint_enc pretrained_ckpts/Cosmos-0.1-Tokenizer-DV4x8x8/encoder.jit \
    --checkpoint_dec pretrained_ckpts/Cosmos-0.1-Tokenizer-DV4x8x8/decoder.jit

代码架构：清晰的模块化设计

项目核心代码位于 cosmos_tokenizer/ 目录，结构如下：

cosmos_tokenizer/
├── video_lib.py        # 视频分词器核心类 CausalVideoTokenizer
├── image_lib.py        # 图像分词器核心类 ImageTokenizer
├── video_cli.py        # 视频分词 CLI 入口
├── image_cli.py        # 图像分词 CLI 入口
├── utils.py            # 模型加载、格式转换等工具函数
├── networks/           # 网络架构定义
│   ├── continuous_video.py   # 连续视频自动编码器（因果 3D 卷积）
│   ├── continuous_image.py   # 连续图像自动编码器
│   ├── discrete_video.py     # 离散视频（VQ/FSQ 量化）
│   └── discrete_image.py     # 离散图像
└── modules/            # 基础组件
    ├── layers2d.py     # 2D 卷积、归一化层
    ├── layers3d.py     # 3D 卷积（时序卷积）
    ├── quantizers.py   # VQ、FSQ 量化器实现
    ├── distributions.py # 概率分布层
    ├── patching.py     # patch embedding
    └── utils.py        # 小工具

核心依赖极为精简：loguru（日志）、mediapy（媒体读写）、einops/einx（张量重塑）、huggingface-hub（权重下载）。值得注意的是，Cosmos Tokenizer 不需要 Diffusers 或 Transformers 这类重型框架——它是一个独立的推理库，可以直接嵌入任何 pipeline。

图：Cosmos Tokenizer 在 A100 80GB GPU 上的推理速度对比，最多比 SOTA 快 12 倍

上手体验：从安装到跑通 demo 全流程

方式一：Colab notebook（推荐）

项目提供了两个可直接在 Google Colab 打开的 notebook：

• Image Tokenization
• Video Tokenization

Colab 环境已预装 PyTorch 和 CUDA，只需安装 pip install -e . 并从 Hugging Face 下载权重即可运行，全程无需配置本地环境。

方式二：Docker 部署

git clone https://github.com/NVIDIA/Cosmos-Tokenizer.git
cd Cosmos-Tokenizer
docker build -t cosmos-tokenizer -f Dockerfile .
docker run --gpus all -it --rm -v /home/${USER}:/home/${USER} \
    --workdir ${PWD} cosmos-tokenizer /bin/bash

Dockerfile 基于 pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel，基础镜像已包含 CUDA 12.1 和 cuDNN 9，开箱即用。唯一额外依赖是 ffmpeg 和 git-lfs（用于下载测试视频）。

方式三：pip 安装

apt-get install -y ffmpeg git-lfs
git clone https://github.com/NVIDIA/Cosmos-Tokenizer.git
cd Cosmos-Tokenizer
git lfs pull
pip3 install -e .

安装后可直接调用 cosmos-image 和 cosmos-video 两个 CLI 命令。

门槛评估

维度	评价
Python 技能要求	中等（需了解 PyTorch tensor 格式）
硬件要求	必须 NVIDIA GPU + CUDA 12.1
最小显存	8GB（连续版 CV4x8x8）
安装复杂度	低（pip install 或 Docker）
文档完整性	高（README 详尽，含 Colab）

最大门槛：必须使用 NVIDIA GPU。没有 NVIDIA 显卡的话，目前无法在 CPU 或 Apple Silicon 上运行——项目未提供 ONNX 或其他跨平台推理方案。

性能与局限：不吹不黑

性能亮点

在 NVIDIA 官方 benchmark 中，Cosmos Tokenizer 在 DAVIS 数据集上的压缩-重建质量（PSNR/SSIM）和速度均大幅领先此前的 SOTA 方法：

• 压缩率：最高 2048 倍（8×16×16），比主流方案多压缩 8 倍
• 重建质量：8×8×8 压缩率下，PSNR 超过 30dB（肉眼几乎无法分辨）
• 推理速度：A100 80GB 上，单帧编码+解码耗时比竞争对手快 2—12 倍

现存局限

1. CUDA Only：最核心的限制。没有 AMD ROCm、CPU 或 Apple Metal 支持，对非 NVIDIA 用户不友好。
2. 预训练模型需手动下载：Hugging Face 权重约数 GB，README 只提供了手动下载脚本，没有集成 huggingface-cli 自动下载的完整引导。
3. 仓库已归档：2025 年 2 月起只读，不再接受 PR，新模型和功能已迁移至 Cosmos 主仓库。
4. TensorRT 尚未就绪：README 提到 TensorRT 加速"即将到来"，但截至归档时尚未集成。
5. 仅推理，无训练代码：项目只发布了推理权重和推理代码，训练 pipeline 属于 Cosmos 商业版的一部分。

行业意义：视频生成民主化的关键一环

Cosmos Tokenizer 的出现标志着视频 tokenization 从"各自为政"走向"统一基础设施"。此前，每个视频生成项目几乎都要自己训练 tokenizer，数据量小的团队往往只能用质量一般的通用 tokenizer。NVIDIA 一次性开源 12 款覆盖不同压缩率的 tokenizer，让中小团队也能用上顶级的视频压缩工具。

从生态角度看，Cosmos Tokenizer 是 NVIDIA Cosmos 全栈 AI 战略的一部分。底层是 tokenization，中间是 Cosmos 世界基础模型，上层是物理 AI 应用（如机器人、自动驾驶仿真）。Tokenization 就像"视频版的 JPEG 压缩"，但压缩比和画质都远超 JPEG，是现代视频生成管线中不可或缺的"水电煤"。

未来，随着 TensorRT 优化和更多硬件后端的支持，Cosmos Tokenizer 有望成为视频 AI 领域的标配工具。