NVIDIA Cosmos Tokenizer 完全指南:图像与视频神经分词器(Neural Tokenizer)的原理、安装与推理实战

📅 发布时间:2026/9/15 19:24:59
NVIDIA Cosmos Tokenizer 完全指南:图像与视频神经分词器(Neural Tokenizer)的原理、安装与推理实战
NVIDIA Cosmos Tokenizer 完全指南图像与视频神经分词器Neural Tokenizer的原理、安装与推理实战【免费下载链接】cosmosNVIDIA Cosmos is an open platform of world models, datasets, and tools that enables developers to build Physical AI for robots, autonomous vehicles, smart infrastructure, and more.项目地址: https://gitcode.com/GitHub_Trending/cosmos7/cosmos导读本文以 cosmos1/models/tokenizer/README.md 为主体系统讲解 NVIDIA Cosmos 平台的核心组件Cosmos Tokenizer——一套面向图像与视频的神经分词器套件。它既能把视觉输入压缩为连续潜空间Continuous Latents也能量化为离散词元Discrete Tokens是构建大规模自回归 Transformer如 LLM 式视频模型与扩散生成器的关键基础设施。读完本文你将掌握其四象限模型选型CI/DI/CV/DV、空间 8x/16x 与时间 4x/8x 的压缩机制、JIT 与原生 PyTorch 两种推理模式、CLI 与 Python API 的完整调用方式并理解 Haar 小波 patching 与 FSQ 有限标量量化等底层实现原理。一、什么是 Cosmos Tokenizer从像素到词元的桥梁Cosmos Tokenizer 是 NVIDIA Cosmos 世界模型平台面向 Physical AI服务于机器人、自动驾驶、智慧基础设施等场景的核心基础组件。它的职责非常纯粹给定一张图像或一段视频输出连续的潜在表示latents或离散的 token 序列。这套 tokenizer 的定位正如 README 所述是推动大规模自回归 Transformer如 LLM 类模型或扩散生成器走向可扩展、稳健、高效开发的基石。对开发者而言它的价值在于将高维、冗余的原始像素转化为紧凑、规则、便于序列模型建模的表示从而让后续的世界模型训练与推理成为可能。Cosmos Tokenizer 共包含4 个模型族覆盖「图像/视频」×「连续/离散」两个维度连续Continuous, C离散Discrete, D图像Images, ICosmos-Tokenizer-CICosmos-Tokenizer-DI视频Videos, VCosmos-Tokenizer-CVCosmos-Tokenizer-DV其整体架构如下——输入视频先经3D Haar Wavelet处理进入Causal Encoder含时间维度的因果卷积与因果注意力分别产出连续与离散两种潜在空间再由Causal Decoder解码最终通过 3D Haar Wavelet 逆变换重建视频核心能力最高 2048x 的总压缩率README 明确指出 Cosmos Tokenizer 的关键量化指标空间压缩率8x 或 16x时间压缩率4x 或 8x总压缩率最高可达 2048x 8x × 16x × 16x。相比此前的最优方法SOTAREADME 声称 Cosmos Tokenizer 在总压缩率上高出 8 倍同时保持更优的图像质量且推理速度快 2x12x。这些是官方文档声明的项目事实其具体证据可见后文「评估与性能」章节引用的官方对比图以及仓库中的推理实现inference/video_lib.py、inference/image_lib.py。二、环境安装Docker 一键就绪Cosmos Tokenizer 的推理代码需要 GPU 环境默认cuda、bfloat16精度。官方推荐使用仓库自带的 Docker 镜像所有命令均在容器内执行。完整步骤见根目录 INSTALL.md安装 NVIDIA Container Toolkit安装指南仅支持 Ubuntu 24.04 / 22.04 / 20.04克隆仓库并构建、启动容器git clone gitgithub.com:NVIDIA/Cosmos.git cd Cosmos docker build -t cosmos . docker run -d --name cosmos_container --gpus all --ipchost -it -v $(pwd):/workspace cosmos docker attach cosmos_container构建完成后容器内的python3环境已包含torch、mediapy等推理所需依赖可直接进入「下载模型」与「运行推理」环节。三、下载预训练权重12 个模型一张表NVIDIA 在 Hugging Face 上托管了全部12 个 Cosmos-Tokenizer 模型覆盖 CI/DI/CV/DV 四族在不同压缩率下的组合。模型命名规则为Cosmos-{版本}-Tokenizer-{类型}{压缩率}模型名称类型压缩率时间×高×宽Cosmos-0.1-Tokenizer-CI8x8连续图像8x8Cosmos-0.1-Tokenizer-CI16x16连续图像16x16Cosmos-0.1-Tokenizer-CV4x8x8连续视频4x8x8Cosmos-0.1-Tokenizer-CV8x8x8连续视频8x8x8Cosmos-0.1-Tokenizer-CV8x16x16连续视频8x16x16Cosmos-0.1-Tokenizer-DI8x8离散图像8x8Cosmos-0.1-Tokenizer-DI16x16离散图像16x16Cosmos-0.1-Tokenizer-DV4x8x8离散视频4x8x8Cosmos-0.1-Tokenizer-DV8x8x8离散视频8x8x8Cosmos-0.1-Tokenizer-DV8x16x16离散视频8x16x16Cosmos-1.0-Tokenizer-CV8x8x8连续视频1.0 版8x8x8Cosmos-1.0-Tokenizer-DV8x16x16离散视频1.0 版8x16x16使用下面的 Python 片段即可批量下载全部权重需先在 Hugging Face 生成 Access Tokenfrom huggingface_hub import login, snapshot_download import os login(tokenYOUR-HF-TOKEN, add_to_git_credentialTrue) model_names [ Cosmos-0.1-Tokenizer-CI8x8, Cosmos-0.1-Tokenizer-CI16x16, Cosmos-0.1-Tokenizer-CV4x8x8, Cosmos-0.1-Tokenizer-CV8x8x8, Cosmos-0.1-Tokenizer-CV8x16x16, Cosmos-0.1-Tokenizer-DI8x8, Cosmos-0.1-Tokenizer-DI16x16, Cosmos-0.1-Tokenizer-DV4x8x8, Cosmos-0.1-Tokenizer-DV8x8x8, Cosmos-0.1-Tokenizer-DV8x16x16, Cosmos-1.0-Tokenizer-CV8x8x8, Cosmos-1.0-Tokenizer-DV8x16x16, ] for model_name in model_names: hf_repo nvidia/ model_name local_dir checkpoints/ model_name print(fdownloading {model_name}...) snapshot_download(repo_idhf_repo, local_dirlocal_dir)每个 checkpoint 目录checkpoints/{model_name}/下提供三份TorchScriptJIT模型├── Cosmos-1.0-Tokenizer-CV8x8x8/ │ ├── encoder.jit │ ├── decoder.jit │ └── autoencoder.jitencoder.jit仅编码器像素 → 潜空间 / tokendecoder.jit仅解码器潜空间 / token → 像素autoencoder.jit完整的编码-解码一体模型可直接做端到端重建。说明snapshot_download会下载仓库内全部文件其中 JIT 模型即上表三件套。下载脚本见仓库根目录 scripts/download_autoregressive.py 的同款模式。四、命令行推理图像/视频的编码与重建仓库为图像与视频分别提供了 CLI 入口支持glob通配符批量处理。同一套命令对连续/离散分词器都适用只需把--checkpoint_enc/--checkpoint_dec或整体--checkpoint指向对应的 JIT 权重。4.1 图像自动编码Autoencoding以Cosmos-CI8x8 压缩为例对测试图片cosmos1/models/tokenizer/test_data/image.png做「编码 → 解码」重建model_nameCosmos-0.1-Tokenizer-CI8x8 python3 -m cosmos1.models.tokenizer.inference.image_cli \ --image_pattern cosmos1/models/tokenizer/test_data/image.png \ --checkpoint_enc checkpoints/${model_name}/encoder.jit \ --checkpoint_dec checkpoints/${model_name}/decoder.jit若未指定--output_dir重建结果默认写到输入目录下的reconstructions子目录即cosmos1/models/tokenizer/test_data/reconstructions/image.png该默认行为由 inference/utils.py 中的get_output_filepath实现。4.2 视频自动编码以Cosmos-DV4x8x8 压缩为例对测试视频cosmos1/models/tokenizer/test_data/video.mp4做重建model_nameCosmos-0.1-Tokenizer-DV4x8x8 python3 -m cosmos1.models.tokenizer.inference.video_cli \ --video_pattern cosmos1/models/tokenizer/test_data/video.mp4 \ --checkpoint_enc checkpoints/${model_name}/encoder.jit \ --checkpoint_dec checkpoints/${model_name}/decoder.jit默认输出路径为cosmos1/models/tokenizer/test_data/reconstructions/video.mp4。4.3 CLI 参数速查以 video_cli.py 为准参数取值默认值说明--video_pattern/--image_patternglob 路径无输入文件通配符支持批量--checkpoint路径None完整 autoencoder JIT 模型--checkpoint_enc路径None编码器 JIT 模型--checkpoint_dec路径None解码器 JIT 模型--modejit/torchjit推理后端见第五章--tokenizer_type图像CI/DI视频CV/DV无仅torch模式需要--spatial_compression8 或 168空间压缩因子torch模式用--temporal_compression4 或 84时间压缩因子torch模式用仅视频--short_sizeintNone将短边缩放至该尺寸后再推理--temporal_windowint17视频按该帧数窗口滑窗处理仅视频--dtype如bfloat16bfloat16推理精度--device如cudacuda运行设备--output_dir路径None输出目录默认输入目录/reconstructions--output_fpsfloat24.0输出视频帧率仅视频--save_inputflag关闭同时保存输入副本_input后缀注意当--checkpoint_enc、--checkpoint_dec、--checkpoint三者都未提供时CLI 会中止并提示需要编码器/解码器或完整 autoencoder见 image_cli.py。五、两种推理后端JIT 与原生 PyTorch默认情况下 CLI 加载encoder.jit/decoder.jit这类TorchScript 编译模型直接推理无需任何网络定义代码。若要运行原生 PyTorch 模式在命令后追加--modetorch即可PyTorch 模式下模型会根据仓库内的原生网络定义脚本networks/重新实例化因此必须额外提供参数来指定网络结构权重仍来自 JIT 文件--checkpoint_enc、--checkpoint_dec、--checkpoint依旧指向 JIT 路径加载时会从 JIT 模型中提取state_dict来初始化原生 PyTorch 模型。5.1 图像原生 PyTorch 实例化Cosmos-DI空间压缩 8x需要追加的三个参数--modetorch--tokenizer_typeDI--spatial_compression8model_nameCosmos-0.1-Tokenizer-DI8x8 python3 -m cosmos1.models.tokenizer.inference.image_cli \ --image_pattern cosmos1/models/tokenizer/test_data/*.png \ --modetorch \ --tokenizer_typeDI \ --spatial_compression8 \ --checkpoint_enc checkpoints/${model_name}/encoder.jit \ --checkpoint_dec checkpoints/${model_name}/decoder.jit5.2 视频原生 PyTorch 实例化Cosmos-CV8x8x8需要追加四个参数--modetorch--tokenizer_typeCV--temporal_compression8--spatial_compression8model_nameCosmos-1.0-Tokenizer-CV8x8x8 python3 -m cosmos1.models.tokenizer.inference.video_cli \ --video_pattern cosmos1/models/tokenizer/test_data/*.mp4 \ --modetorch \ --tokenizer_typeCV \ --temporal_compression8 \ --spatial_compression8 \ --checkpoint_enc checkpoints/${model_name}/encoder.jit \ --checkpoint_dec checkpoints/${model_name}/decoder.jit5.3 PyTorch 模式背后的实现从源码看torch模式的实质是CLI 根据tokenizer_type从 networks/init.py 的TokenizerConfigs枚举中取出默认配置字典再用命令行参数覆盖spatial_compression/temporal_compression随后 inference/utils.py 中的_load_pytorch_model依据name字段从TokenizerModels枚举CI→ContinuousImageTokenizer、DI→DiscreteImageTokenizer、CV→CausalContinuousVideoTokenizer、DV→CausalDiscreteVideoTokenizer实例化网络最后把 JIT 权重load_state_dict(strictFalse)灌入。这意味着你完全可以用这套配置机制构造自己的自定义压缩率网络。六、Python API 实战编码为连续潜空间 / 离散 token除 CLI 外仓库提供面向库调用的 Python 类图像用ImageTokenizerinference/image_lib.py视频用CausalVideoTokenizerinference/video_lib.py。两者均以torch.nn.Module形式封装支持传入完整 autoencoder 或独立的 encoder/decoder JIT 路径默认运行在cuda、bfloat16。6.1 编码为连续潜空间Continuous Latentsimport torch from cosmos1.models.tokenizer.inference.video_lib import CausalVideoTokenizer model_name Cosmos-0.1-Tokenizer-CV4x8x8 input_tensor torch.randn(1, 3, 9, 512, 512).to(cuda).to(torch.bfloat16) # [B, C, T, H, W] encoder CausalVideoTokenizer(checkpoint_encfcheckpoints/{model_name}/encoder.jit) (latent,) encoder.encode(input_tensor) torch.testing.assert_close(latent.shape, (1, 16, 3, 64, 64)) # 解码重建 decoder CausalVideoTokenizer(checkpoint_decfcheckpoints/{model_name}/decoder.jit) reconstructed_tensor decoder.decode(latent) torch.testing.assert_close(reconstructed_tensor.shape, input_tensor.shape)这里的latent形状为(1, 16, 3, 64, 64)语义拆解如下输入[B1, C3, T9, H512, W512]9 帧 512x512 RGB 视频时间维度 9 → 3时间压缩 4x 对应1 (T-1)/4 3个时间潜位置第一个潜位置代表首帧即因果式对齐源码注释见 video_lib.py空间维度 512 → 648x8 空间压缩C16连续潜空间的通道数该常量在 networks/configs.py 的latent_channels16/z_channels16中定义。6.2 编码为离散 tokenDiscrete Tokensimport torch from cosmos1.models.tokenizer.inference.video_lib import CausalVideoTokenizer model_name Cosmos-0.1-Tokenizer-DV4x8x8 input_tensor torch.randn(1, 3, 9, 512, 512).to(cuda).to(torch.bfloat16) # [B, C, T, H, W] encoder CausalVideoTokenizer(checkpoint_encfcheckpoints/{model_name}/encoder.jit) (indices, codes) encoder.encode(input_tensor) torch.testing.assert_close(indices.shape, (1, 3, 64, 64)) torch.testing.assert_close(codes.shape, (1, 6, 3, 64, 64)) # 解码重建 decoder CausalVideoTokenizer(checkpoint_decfcheckpoints/{model_name}/decoder.jit) reconstructed_tensor decoder.decode(indices) torch.testing.assert_close(reconstructed_tensor.shape, input_tensor.shape)离散模式的输出有两个indices形状(1, 3, 64, 64)元素为[1..64K]区间的整数值——64K65536即 FSQ 隐式码本大小由 levels(8,8,8,5,5,5)相乘得到8×8×8×5×5×5 64000README 以 64K 约称。第一个整数图对应首帧codes量化前的连续潜变量形状(1, 6, 3, 64, 64)其中C6是FSQ 的 level 维度数6 个标量量化通道。6.3 端到端重建与任意长度视频若构造时传入checkpoint完整 autoencoder JITautoencode()会直接跑完整模型否则内部自动串联encode → decode见 video_lib.py。forward(video, temporal_window17)支持任意长度视频的滑动窗口重建按temporal_window逐窗处理先空间补零 时间边缘反射 padding 到对齐尺寸推理后再裁剪回原尺寸并拼接实现见 inference/utils.py 的pad_video_batch/unpad_video_batch。七、源码级原理压缩率从哪来7.1 Haar 小波 patching无参数下采样图像/视频在进入主干网络前先经过Patcher做可逆下采样modules/patching.py2D 图像使用 Haar 小波_WAVELETS {haar: [0.7071, 0.7071], ...}每做一次离散小波变换DWT把通道从c扩到4c、空间尺寸减半——例如[3, 256, 256] → [12, 128, 128] → [48, 64, 64]源码 docstring 中的示例3D 视频使用Patcher3D每次 DWT 将通道扩到 8 倍、三个维度各减半同时对首帧做 repeat_interleave以维持因果对齐对应的UnPatcher/UnPatcher3D用转置卷积完成逆小波变换恢复像素。patch_size4时叠加两次 DWT正好贡献 4x 的空间下采样。这正是「patch_methodhaar」配置见 networks/configs.py 中各 dict 的patch_size4, patch_methodhaar的含义——小波变换完全无学习参数比可学习的 stride 卷积更高效且可逆。7.2 FSQ 有限标量量化无需码本查找的离散化离散 tokenizerDI/DV使用FSQFinite Scalar Quantization而非传统 VQ 码本查找modules/quantizers.py每个潜在向量的 6 个通道分别按levels[8,8,8,5,5,5]的整数级数独立量化bound中用 tanh 约束范围round_ste保留直通梯度隐式码本大小 各级数之积 64000约 64K索引由codes_to_indices用混合进制_basis计算得到解码时由indices_to_codes反向还原由于无需查询嵌入表FSQ 天然支持超大批次与高吞吐这也是官方宣称推理更快的原因之一。7.3 因果视频 tokenizer首帧先行的时序设计视频族CV/DV的关键在于Causal因果设计时间维度压缩通过因果卷积与因果注意力实现保证解码时任一帧只依赖其之前含自身的帧从而支持流式生成时间压缩 T→t 满足t 1 (T-1)/压缩率例如 9 帧 4x 时间压缩 → 3 个时间潜位置首帧潜位置与首帧严格对齐README 与 video_lib.py 的 docstring 均明确此约定。7.4 默认网络配置速查networks/configs.py 中四个族的关键超参如下配置项CIDICVDVchannels基础通道数128128128128channels_mult各分辨率通道倍增[2,4,4][2,4,4][2,4,4][2,4,4]num_res_blocks2222patch_size/patch_method4 / haar4 / haar4 / haar4 / haar默认空间压缩1616816默认时间压缩——88z_channels/latent_channels16 / 16256 / —16 / 1616 / —quantizer—FSQ—FSQembedding_dim/levels—6 / [8,8,8,5,5,5]—6 / [8,8,8,5,5,5]提示表中为各族的默认压缩率CLI 的--spatial_compression/--temporal_compression可在 PyTorch 模式下覆盖默认值但必须与所下载权重的实际压缩率一致否则加载的state_dict无法匹配。八、NeMo 集成JIT 推理与数据集 tokenization除了仓库自带的推理代码Cosmos Tokenizer 还深度集成到 NVIDIA NeMo 框架nemo/README.md支持JIT 推理安装 NeMoGitHub main 分支后直接用CausalVideoTokenizer.from_pretrained(model_name)一行加载模型import torch from nemo.collections.common.video_tokenizers.cosmos_vision_tokenizer import CausalVideoTokenizer model_name Cosmos-0.1-Tokenizer-CV4x8x8 model CausalVideoTokenizer.from_pretrained(model_name) input_tensor torch.randn(1, 3, 9, 512, 512).to(cuda).to(torch.bfloat16) (latent, ) model.encode(input_tensor)数据集 tokenization 与多模态训练NeMo 提供了用 Cosmos Tokenizer 构造多模态训练数据集的示例可直接对接后续的世界模型训练流程NeMo 后训练Post-trainingNeMo 框架目前支持对Cosmos-1.0-Tokenizer-CV8x8x8与Cosmos-1.0-Tokenizer-DV8x16x16两个模型做领域微调尤其针对自动驾驶数据推荐在 H100-80GB 或 A100-80GB 上运行具体步骤见 nemo/README.md。九、评估与性能官方基准结果9.1 DAVIS 数据集量化对比README 报告了在 DAVISPerazzi et al., 2016数据集上与历史 tokenizer 的定量对比即使在更高的压缩率8x8x8 与 8x16x16下Cosmos Tokenizer 仍优于此前方法展现了出色的「压缩-质量」权衡9.2 参数量与推理耗时README 还对比了各模型在单张 A100 80GB GPU 上的参数量与每图/每帧平均编码、解码耗时Cosmos Tokenizer 相比此前方法快 2x12x且模型尺寸更小体现了极高的 tokenization 效率以上数据均为 README 官方声明复现与进一步验证可参考官方公开的 TokenBench 基准见下节及仓库内的评估代码。十、TokenBench视频 tokenizer 的标准化评测基准为标准化 Cosmos Tokenizer 的评估官方还配套发布了TokenBench——一个覆盖机器人操作、驾驶、第一人称视角egocentric、网络视频等多领域的综合基准包含高分辨率、长时长视频专为评测视频 tokenizer 设计。TokenBench 面向社区公开可用于与其他视频 tokenizer 在同一标准下横向比较压缩质量与效率。十一、许可证与引用模型权重采用 NVIDIA Open Model License该协议确认模型可商用允许创建和分发衍生模型NVIDIA 不对使用模型或衍生模型生成的任何输出主张所有权。仓库代码本仓库含 tokenizer 代码采用Apache 2.0许可见根目录 LICENSE。若在你的工作中使用了 Cosmos Tokenizer请按如下方式引用article{agarwal2025cosmos, title{Cosmos World Foundation Model Platform for Physical AI}, author{NVIDIA et. al.}, journal{arXiv preprint arXiv:2501.03575}, year{2025} }十二、致谢与源码归属仓库注明modules/ 目录中的部分代码实现派生自以下开源项目CompVis/stable-diffusion、lucidrains/magvit2-pytorch、lucidrains/vector-quantize-pytorch、CompVis/taming-transformers。这也解释了为什么 modules/quantizers.py 中 FSQ/LFQ/VQ 的实现与上述项目保持接口与算法的一致性——理解这些上游约定有助于你在二次开发时快速定位量化器行为。总结Cosmos Tokenizer 以「连续/离散 × 图像/视频」四象限覆盖视觉 tokenization 的全部主流需求通过 Haar 小波 patching、因果式 3D 网络与 FSQ 量化把最高 2048x 的总压缩率、64K 隐式码本与高效推理组合在一起。无论是想快速用 CLI 做图像/视频重建、用 Python API 提取潜空间或离散 token还是打算对接 NeMo 做数据集构建与领域后训练上文给出的模型选型表、参数速查表与源码级原理分析都已覆盖下一步即可按 INSTALL.md 搭建环境、下载对应 checkpoint 并跑通第一条推理命令。【免费下载链接】cosmosNVIDIA Cosmos is an open platform of world models, datasets, and tools that enables developers to build Physical AI for robots, autonomous vehicles, smart infrastructure, and more.项目地址: https://gitcode.com/GitHub_Trending/cosmos7/cosmos创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考