File size: 22,580 Bytes
3f20bfe 68b3596 3f20bfe 197063c 3f20bfe 259e7f0 3f20bfe 3ae1b01 3f20bfe |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 |
---
license: other
license_name: nvidia-open-model-license
license_link: >-
https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf
library_name: nemo
---
# **Cosmos Tokenizer**: A suite of image and video tokenizers
[**Website**](https://research.nvidia.com/labs/dir/cosmos-tokenizer) | [**Code**](https://github.com/NVIDIA/Cosmos-Tokenizer) | [**Video**](https://youtu.be/Soy_myOfWIU)
# Model Overview
## Description:
**Cosmos Tokenizer** is a suite of visual tokenizers for images and videos that delivers various compression rates while maintaining high reconstruction quality. Cosmos Tokenizer can serve as an effective and efficient building block in both diffusion-based and autoregressive models for image and video generation.
Our tokenizers come in two types: **Continuous** (C) and **Discrete** (D), each with **Image** (I) and **Video** (V) variants:
* Continuous tokenizers encode visual data into continuous latent embeddings, as shown in latent diffusion models like [Stable Diffusion](https://github.com/CompVis/stable-diffusion). These embeddings are suitable for models that generate data by sampling from continuous distributions.
* Discrete tokenizers encode visual data into discrete latent codes, mapping them into quantized indices, as seen in autoregressive transformers such as [VideoPoet](https://sites.research.google/videopoet/). This discretization is required for models that generate data by optimizing the cross-entropy loss, such as the GPT models.
| | Continuous ( C ) | Discrete ( D ) |
| ------------------|---------------------|---------------------|
| **Images ( I )** | Cosmos-Tokenizer-CI | Cosmos-Tokenizer-DI |
| **Videos ( V )** | Cosmos-Tokenizer-CV | Cosmos-Tokenizer-DV |
Given an image or a video, Cosmos Tokenizer outputs either continuous latents or discrete tokens. Cosmos Tokenizer achieves spatial compression rates of 8x8 or 16x16 and temporal compression factors of 4x or 8x, resulting in a total compression factor of up to 2048x (=8x16x16). Cosmos Tokenizer delivers 8x more total compression than state-of-the-art (SOTA) methods while simultaneously maintaining higher image quality and running up to 12x faster than the best available SOTA tokenizers.
**Model Developer**: NVIDIA
## Model Versions
The initial release (v1.0) of Cosmos Tokenizer includes the following tokenizers:
* **Continuous Tokenizers**
* Continuous Image (CI) Tokenizer
* [Cosmos-Tokenizer-CI8x8](https://huggingface.co/nvidia/Cosmos-Tokenizer-CI8x8) (8x8 spatial compression)
* [Cosmos-Tokenizer-CI16x16](https://huggingface.co/nvidia/Cosmos-Tokenizer-CI16x16) (16x16 spatial compression)
* Continuous Video (CV) Tokenizer
* [Cosmos-Tokenizer-CV4x8x8](https://huggingface.co/nvidia/Cosmos-Tokenizer-CV4x8x8) (4x temporal compression, 8x8 spatial compression)
* [Cosmos-Tokenizer-CV8x8x8](https://huggingface.co/nvidia/Cosmos-Tokenizer-CV8x8x8) (8x temporal compression, 8x8 spatial compression)
* [Cosmos-Tokenizer-CV8x16x16](https://huggingface.co/nvidia/Cosmos-Tokenizer-CV8x16x16) (8x temporal compression, 16x16 spatial compression)
* **Discrete Tokenizers**
* Discrete Image (DI) Tokenizer
* [Cosmos-Tokenizer-DI8x8](https://huggingface.co/nvidia/Cosmos-Tokenizer-DI8x8) (8x8 spatial compression)
* [Cosmos-Tokenizer-DI16x16](https://huggingface.co/nvidia/Cosmos-Tokenizer-DI16x16) (16x16 spatial compression)
* Discrete Video (DV) Tokenizer
* [Cosmos-Tokenizer-DV4x8x8](https://huggingface.co/nvidia/Cosmos-Tokenizer-DV4x8x8) (4x temporal compression, 8x8 spatial compression)
* [Cosmos-Tokenizer-DV8x8x8](https://huggingface.co/nvidia/Cosmos-Tokenizer-DV8x8x8) (8x temporal compression, 8x8 spatial compression)
* [Cosmos-Tokenizer-DV8x16x16](https://huggingface.co/nvidia/Cosmos-Tokenizer-DV8x16x16) (8x temporal compression, 16x16 spatial compression)
### License/Terms of Use:
[NVIDIA Open Model License](https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf)
Under the NVIDIA Open Model License, NVIDIA confirms:
* Models are commercially usable.
* You are free to create and distribute Derivative Models.
* NVIDIA does not claim ownership to any outputs generated using the Models or Derivative Models.
## Model Architecture:
We designed Cosmos Tokenizer using a lightweight and computationally efficient architecture, featuring a temporally causal design. Specifically, we employ causal temporal convolution and causal temporal attention layers to preserve the natural temporal order of video frames, ensuring seamless tokenization of images and videos using a single unified network architecture. The encoder and decoder form a symmetrical pair, which are mirrors of each other. The encoder starts with a 2-level [Haar wavelet](https://link.springer.com/book/10.1007/978-3-319-04295-4) transform layer, which down-samples inputs by a factor of 4 in both spatial and temporal dimensions. Likewise, the decoder ends with an inverse wavelet transform. We employ the vanilla autoencoder (AE) formulation to model the latent space for continuous tokenizers. For discrete tokenizers, we adopt the [Finite-Scalar-Quantization](https://openreview.net/forum?id=8ishA3LxN8) (FSQ) as the latent space quantizer.
## Input/Output Specifications
### Encoder
* **Input**
* **Types:** Images or Videos
* **Format:** RGB (Red, Green, Blue)
* **Resolution:**
* Minimum: 256px (shorter side)
* Maximum: Up to 4K
* **Video Length:** Up to 8 seconds for 1080p videos (bounded by A100 80G GPU memory; higher resolutions will have shorter supported durations)
* **Output**
* **Types:** Tokens
* Continuous Image/Video Tokenizers: Continuous value feature vectors
* Discrete Image/Video Tokenizers: Integer indices
### Decoder
* **Input**
* **Types:** Tokens from encoder
* **Output**
* **Types:** Images or Videos (matching input type)
* **Format:** RGB (Red, Green, Blue)
* **Resolution:** Same as input resolution
* **Video Length:** Same as input video length
## Software Integration (Required For NVIDIA Models Only):
**Runtime Engine(s):**
* [Cosmos-Tokenizer](https://github.com/NVIDIA/Cosmos-Tokenizer)
* [NeMo](https://github.com/NVIDIA/NeMo) (please install the latest version from the GitHub main branch)
**Supported Hardware Microarchitecture Compatibility:**
* NVIDIA Ampere (e.g., A100)
* NVIDIA Hopper (e.g., H100)
Note: We have only tested Cosmos Tokenizer with BF16 precision on Ampere and Hopper GPUs. If you are using older versions of NVIDIA GPUs (e.g., NVIDIA Volta GPUs), you may need to switch to FP32 precision.
**Operating System(s):**
* Linux (We have not tested on other operating systems.)
# Usage
Inference Engines:
* [Cosmos-Tokenizer](https://github.com/NVIDIA/Cosmos-Tokenizer) (PyTorch)
* [NeMo](https://github.com/NVIDIA/NeMo)
## Inference with `Cosmos-Tokenizer` (PyTorch)
### Step-1: Installation of `Cosmos-Tokenizer`
Note: Currently, the `Cosmos-Tokenizer` code is only supported on Linux.
- Please clone the `Cosmos-Tokenizer` from GitHub repo [github.com/NVIDIA/Cosmos-Tokenizer](https://github.com/NVIDIA/Cosmos-Tokenizer).
```bash
git clone https://github.com/NVIDIA/Cosmos-Tokenizer.git
cd Cosmos-Tokenizer
```
- Install dependencies
```bash
pip3 install -r requirements.txt
apt-get install -y ffmpeg
```
- Preferably, you could build a docker image using our provided Dockerfile.
```bash
docker build -t cosmos-docker -f Dockerfile.
# You can run the container as:
docker run --gpus all -it --rm -v /home/${USER}:/home/${USER} \
--workdir ${PWD} cosmos-docker /bin/bash
```
### Step-2: Download Pre-trained Checkpoints
- Create a local directory for the pre-trained checkpoints and download the
pre-trained checkpoints from HuggingFace.
```python
from huggingface_hub import login, snapshot_download
import os
# You could get your Hugging Face token from https://huggingface.co/settings/tokens
login(token=<YOUT-HF-TOKEN>, add_to_git_credential=True)
# You could specify the tokenizers you want to download.
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",
]
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} to {local_dir}...")
snapshot_download(repo_id=hf_repo, local_dir=local_dir)
```
- Under the ech checkpoint directory `pretrained_ckpts/<model-name>`, we provide the encoder,
decoder and the full autoencoder JIT models.
```bash
├── pretrained_ckpts/
│ ├── Cosmos-Tokenizer-DV8x8x8/
│ │ ├── encoder.jit
│ │ ├── decoder.jit
│ │ ├── autoencoder.jit
│ ...
```
### Step-3: Run Inference
You can use the following example commands to encode and decode images or videos. For each, the same command works for both continuous and discrete tokenization. Simply provide the proper JIT-compiled ckpt to `checkpoint_enc`, `checkpoint_dec`, or the full autoencoder ckpt to `checkpoint`.
```python
import torch
from cosmos_tokenizer.image_lib import ImageTokenizer
model_name = "Cosmos-Tokenizer-DI8x8"
input_tensor = torch.randn(1, 3, 512, 512).to('cuda').to(torch.bfloat16) # [B, C, H, W]
encoder = ImageTokenizer(checkpoint_enc=f'pretrained_ckpts/{model_name}/encoder.jit')
(indices,codes) = encoder.encode(input_tensor)
torch.testing.assert_close(indices.shape, (1, 64, 64))
torch.testing.assert_close(codes.shape, (1, 6, 64, 64))
# The input tensor can be reconstructed by the decoder as:
decoder = ImageTokenizer(checkpoint_dec=f'pretrained_ckpts/{model_name}/decoder.jit')
reconstructed_tensor = decoder.decode(indices)
torch.testing.assert_close(reconstructed_tensor.shape, input_tensor.shape)
```
The `indices` will have the shape `(1, 64, 64)` and contain integral values in the range `[1..64K]`.
The `codes` will contain the pre-quantization continuous latent with shape `(1, 6, 64, 64)`, where `C=6` represents the number of FSQ levels.
**Note**: More inference usage commands, including both TorchScript (JIT) and PyTorch Inference APIs on real images and videos, can be found on our GitHub repository [github.com/NVIDIA/Cosmos-Tokenizer](https://github.com/NVIDIA/Cosmos-Tokenizer).
## Inference with NeMo
### Step-1: Install NeMo
Please install NeMo from the GitHub `main` branch following the instructions [here](https://github.com/NVIDIA/NeMo?tab=readme-ov-file#pip-from-a-source-branch).
### Step-2: Run Inference
Run the following code to tokenize the video:
```python
import torch
from nemo.collections.common.video_tokenizers.cosmos_vision_tokenizer import CausalVideoTokenizer
model_name = "Cosmos-Tokenizer-DI8x8"
model = CausalVideoTokenizer.from_pretrained(model_name)
input_tensor = torch.randn(1, 3, 512, 512).to('cuda').to(torch.bfloat16)
(indices, codes) = model.encode(input_tensor)
```
Please see the [Cosmos Tokenizer README within the NeMo repository](https://github.com/NVIDIA/NeMo/tree/main/nemo/collections/common/video_tokenizers) for additional examples to create training datasets with Cosmos Tokenizer.
# Evaluation
## TokenizationPerformance Comparison
We have extensively evaluated the **Cosmos Tokenizer** suite on various image and video benchmark datasets.
For the evaluation of image tokenizers, we follow prior art to evaluate on MS-COCO 2017, ImageNet-1K, FFHQ, and CelebA-HQ. We use the MS-COCO 2017 validation subset of 5,000 images, ImageNet-1K validation subset of 50,000 images, FFHQ subset of 10,000 images, and CelebA-HQ subset of 14,645 images as image evaluation benchmark.
| Tokenizer | Compression Ratio | Quantization | PSNR (MS-COCO) | SSIM (MS-COCO) | rFID (MS-COCO) | PSNR (ImageNet-1K) | SSIM (ImageNet-1K) | rFID (ImageNet-1K) | PSNR (FFHQ) | SSIM (FFHQ) | rFID (FFHQ) | PSNR (CelebA-HQ) | SSIM (CelebA-HQ) | rFID (CelebA-HQ) |
|----------------------|-------------------|--------------|----------------|----------------|----------------|---------------------|---------------------|---------------------|----------------|----------------|----------------|-------------------|-------------------|-------------------|
| Open-MAGVIT2 | 16×16 | LFQ | 30.06 | 0.502 | 6.649 | 29.62 | 0.398 | 2.701 | 31.77 | 0.774 | 1.994 | 32.36 | 0.844 | 2.865 |
| LlamaGen | 8×8 | VQ | 30.71 | 0.616 | **4.123** | 30.28 | 0.498 | **1.403** | 33.39 | 0.868 | 0.701 | 34.82 | 0.937 | 0.502 |
| LlamaGen | 16×16 | VQ | 29.93 | 0.491 | 6.077 | 29.81 | 0.448 | 1.657 | 31.58 | 0.772 | 1.366 | 32.18 | 0.837 | 1.113 |
| Cosmos-Tokenizer-DI | 8×8 | FSQ | **31.74** | **0.730** | 4.564 | **31.73** | **0.725** | 1.841 | **35.35** | **0.892** | **0.555** | **37.77** | **0.948** | **0.261** |
| Cosmos-Tokenizer-DI | 16×16 | FSQ | 30.74 | 0.591 | 12.252 | 30.69 | 0.582 | 6.529 | 33.17 | 0.808 | 7.663 | 33.86 | 0.854 | 5.953 |
* We compare with the state-of-the-art discrete image tokenizers, [Open-MAGVIT2](https://github.com/TencentARC/Open-MAGVIT2) and [LlamaGen](https://github.com/FoundationVision/LlamaGen).
* Evaluation metrics:
* Peak Signal-to-Noise Ratio (PSNR)
* Structural Similarity (SSIM)
* Reconstruction Fréchet Inception Distance (rFID)
## Runtime Comparison
The following table shows the number of parameters and the averaged encoding and decoding times per image or video frame, measured on a single A100 80GB GPU. For comparison, we also list the parameters and average speeds of prior state-of-the-art tokenizer(s) with the same compression ratio.
| Tokenizer | Resolution | Compression Ratio | Parameters | Time (ms) |
|----------------------|------------|-------------------|------------|-----------|
| LlamaGen | 1024x1024 | 8×8 | 70M | 475 |
| Cosmos-Tokenizer-DI | 1024x1024 | 8×8 | 79M | **64.2** |
Note: We benchmarked the runtime for images under the 8x8 compression and videos under the 4×8×8 compression. Tokenizers with different compression ratios are not included in this comparison.
## Ethical Considerations
NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.
For more detailed information on ethical considerations for this model, please see the subcards of Explainability, Bias, Safety & Security, and Privacy below. Please report security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
### Bias
Field | Response
:---------------------------------------------------------------------------------------------------|:---------------
Participation considerations from adversely impacted groups [protected classes](https://www.senate.ca.gov/content/protected-classes) in model design and testing: | None
Measures taken to mitigate against unwanted bias: | None
### Explainability
Field | Response
:------------------------------------------------------------------------------------------------------|:---------------------------------------------------------------------------------
Intended Application & Domain: | Tokenization of images and videos
Model Type: | Auto-Encoder
Intended Users: | Generative AI developers for image and video generation models
Output: | Images/Videos and Latent Tokens
Describe how the model works: | Compresses and decompresses visual input (image/video).
Technical Limitations: | Due to tokenizer compression limitations, some visual information (such as small text and other structured fine details) may not be reconstructed accurately.
Verified to have met prescribed NVIDIA quality standards: | Yes
Performance Metrics: | Peak Signal-to-Noise Ratio (PSNR), Structural Similarity (SSIM), Reconstruction Fréchet Video Distance (rFVD), Reconstruction Fréchet Inception Distance (rFID), Latency
Potential Known Risks: | Tokenizer's output can parse all forms of input, including what may be considered toxic, offensive, or indecent.
Licensing: | [NVIDIA Open Model License](https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf)
### Privacy
Field | Response
:----------------------------------------------------------------------------------------------------------------------------------|:-----------------------------------------------
Generatable or reverse engineerable personal information? | No
Protected class data used to create this model? | None Known
Was consent obtained for any personal data used? | None Known
How often is dataset reviewed? | Before Release
Is a mechanism in place to honor data subject right of access or deletion of personal data? | Not Applicable
If personal collected for the development of the model, was it collected directly by NVIDIA? | Not Applicable
If personal collected for the development of the model by NVIDIA, do you maintain or have access to disclosures made to data subjects? | Not Applicable
If personal collected for the development of this AI model, was it minimized to only what was required? | Not Applicable
Is there provenance for all datasets used in training? | Yes
Does data labeling (annotation, metadata) comply with privacy laws? | Yes
Is data compliant with data subject requests for data correction or removal, if such a request was made? | Not Applicable
### Safety
Field | Response
:---------------------------------------------------|:----------------------------------
Model Application(s): | Tokenization of images and videos
Describe the life critical impact (if present). | None Known
Use Case Restrictions: | See [NVIDIA Open Model License](https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf)
Model and dataset restrictions: | The Principle of least privilege (PoLP) is applied limiting access for dataset generation and model development. Restrictions enforce dataset access during training, and dataset license constraints adhered to. Model checkpoints are made available on Hugging Face, and may become available on cloud providers' model catalog.
### Plus Plus (++) Promise
We value you, the datasets, the diversity they represent, and what we have been entrusted with. This model and its associated data have been:
* Verified to comply with current applicable disclosure laws, regulations, and industry standards.
* Verified to comply with applicable privacy labeling requirements.
* Annotated to describe the collector/source (NVIDIA or a third-party).
* Characterized for technical limitations.
* Reviewed to ensure proper disclosure is accessible to, maintained for, and in compliance with NVIDIA data subjects and their requests.
* Reviewed before release.
* Tagged for known restrictions and potential safety implications.
# Core Contributors
Fitsum Reda, Jinwei Gu, Xian Liu, Songwei Ge, Ting-Chun Wang, Haoxiang Wang, Ming-Yu Liu |