【亲测免费】 NVIDIA Cosmos Tokenizer 使用教程

1. 项目介绍

Cosmos Tokenizer 是由 NVIDIA 开发的一套图像和视频神经网络标记化工具。它通过先进的视觉标记化技术，为构建大型自回归变压器（如 LLMs）或扩散生成器提供了可扩展、健壮和高效的支持。Cosmos Tokenizer 是 NVIDIA Cosmos 视频基础模型平台的核心组件，旨在帮助物理 AI 开发者更好地、更快地构建他们的物理 AI 系统。

2. 项目快速启动

环境准备

首先，确保你的系统已经安装了以下依赖：

• Python 3.x
• pip
• ffmpeg
• git-lfs

克隆代码

git clone https://github.com/NVIDIA/Cosmos-Tokenizer.git
cd Cosmos-Tokenizer

安装依赖

apt-get install -y ffmpeg git-lfs
git lfs pull
pip3 install -e .

或者，你可以选择构建 Docker 镜像：

docker build -t cosmos-tokenizer -f Dockerfile .

启动 Docker 容器：

docker run --gpus all -it --rm -v /home/${USER}:/home/${USER} --workdir ${PWD} cosmos-tokenizer /bin/bash

下载预训练模型

从 Hugging Face 下载预训练模型：

from huggingface_hub import login, snapshot_download
import os

login(token="<YOUR-HF-TOKEN>", add_to_git_credential=True)

model_names = [
    "Cosmos-0.1-Tokenizer-CI8x8",
    "Cosmos-0.1-Tokenizer-CV4x8x8",
    # ... 其他模型
]

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)

编码和解码图像

import torch
from cosmos_tokenizer.video_lib import CausalVideoTokenizer

model_name = "Cosmos-0.1-Tokenizer-CI8x8"
input_tensor = torch.randn(1, 3, 512, 512).to('cuda').to(torch.bfloat16)

encoder = CausalVideoTokenizer(checkpoint_enc=f'pretrained_ckpts/{model_name}/encoder.jit')
(latent,) = encoder.encode(input_tensor)

decoder = CausalVideoTokenizer(checkpoint_dec=f'pretrained_ckpts/{model_name}/decoder.jit')
reconstructed_tensor = decoder.decode(latent)

3. 应用案例和最佳实践

图像标记化

使用 Cosmos-CI 进行图像标记化的例子：

model_name="Cosmos-0.1-Tokenizer-CI8x8"
python3 -m cosmos_tokenizer.image_cli --image_pattern 'test_data/image.png' --checkpoint_enc "pretrained_ckpts/${model_name}/encoder.jit" --checkpoint_dec "pretrained_ckpts/${model_name}/decoder.jit"

视频标记化

使用 Cosmos-DV 进行视频标记化的例子：

model_name="Cosmos-0.1-Tokenizer-DV4x8x8"
python3 -m cosmos_tokenizer.video_cli --video_pattern 'test_data/video.mp4'

4. 典型生态项目

目前，Cosmos Tokenizer 已被广泛应用于多种场景，包括但不限于视频编辑、视频增强、内容分发网络（CDN）优化等。开发者可以利用 Cosmos Tokenizer 的高效压缩和高质量重构特性，构建更加高效和智能的视频处理解决方案。