Cosmos-Tokenizer模型昇腾推理适配

​ Cosmos Tokenizer是英伟达推出的一套图像和视频的标记器，推动了视觉标记化的最先进技术，为大规模、稳健和高效地开发大型自回归变换器（如LLMs）或扩散生成器铺平了道路。给定一张图像或视频，Cosmos Tokenizer 输出连续潜变量或离散标记。Cosmos Tokenizer 实现了 8 倍或 16 倍的空间压缩率和 4 倍或 8 倍的时间压缩因子，最终总压缩因子高达 2048 倍（=8x16x16）。Cosmos Tokenizer 提供的总压缩比最先进的方法（SOTA）高出 8 倍，同时保持更高的图像质量，并且运行速度比现有最佳 SOTA Tokenizer 快 12 倍。

本文介绍将Cosmos Tokenizer模型适配到昇腾算力环境中，并使用4个图片和视频输入模型进行视图编解码再还原的过程。

实验环境：

• 硬件配置: Atlas 800T A2/910B1
• CANN版本：8.0.RC2
• 软件版本: Python 3.10、 PyTorch 2.1.0、torch-npu 2.1.0.post6

模型信息：

Cosmos Tokenizer: A suite of image and video neural tokenizers.

GitHub：GitHub - NVIDIA/Cosmos-Tokenizer: A suite of image and video neural tokenizers

Hugging Face: https://huggingface.co/collections/nvidia/cosmos-tokenizer-67...

包含10个不同类型和压缩比的模型：

model_names = [

"Cosmos-Tokenizer-CI8x8",

"Cosmos-Tokenizer-CI16x16",

"Cosmos-Tokenizer-CV4x8x8",

"Cosmos-Tokenizer-CV8x8x8",

"Cosmos-Tokenizer-CV8x16x16",

"Cosmos-Tokenizer-DI8x8",

"Cosmos-Tokenizer-DI16x16",

"Cosmos-Tokenizer-DV4x8x8",

"Cosmos-Tokenizer-DV8x8x8",

"Cosmos-Tokenizer-DV8x16x16",

]

模型名称中字母对应代表的意思如下，数字如(4x）8x8代表（时）空压缩比。

​

添加图片注释，不超过 140 字（可选）

迁移适配：

使用Pytorch GPU2Ascend工具迁移，自动替换torch-npu，在线推理。

Pytorch GPU2Ascend工具是非常好用、便捷的昇腾迁移工具，关于该工具的使用，后续将专门讲解，敬请期待。

算子适配：

推理过程中发现一些算子不支持，采用精度调整及转cpu的方式可以跑通。

视图demo推理

选取源码中的测试视频图片（cat.mp4、flower.mov）及自定义视频图片（mouse.png、rose.jpg）作为输入，参考源码测试demo，分别对10个不同版本的模型进行推理测试。推理脚本参考源码：

Autoencoding videos using `Cosmos-CV` with a compression rate of 8x8x8. 
model_name="Cosmos-Tokenizer-CV8x8x8" 
python3 -m cosmos_tokenizer.video_cli \
     --video_pattern 'test_data/*.mp4' \
     --mode=torch \ --tokenizer_type=CV \
     --temporal_compression=8 \
     --spatial_compression=8 \
     --checkpoint_enc pretrained_ckpts/${model_name}/encoder.jit \
     --checkpoint_dec pretrained_ckpts/${model_name}/decoder.jit

经测验，4个视图输入、10个版本模型均能输出还原对应结果，分别率有些差异：

添加图片注释，不超过 140 字（可选）
​