Cosmos Tokenizer: A suite of image and video tokenizers

Website | NVIDIA News | NVIDIA Blog | Hugging Face | YouTube

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. This model is ready for commercial use.

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. 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. This discretization is required for models that generate data by optimizing the cross-entropy loss, such as the GPT models.

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

This release (v1.0) of Cosmos Tokenizer includes the following tokenizers:

• Continuous Tokenizers
  • Cosmos-1.0-Tokenizer-CV8x8x8 (8x temporal compression, 8x8 spatial compression, 720 short spatial size, 121 frames context)
• Discrete Tokenizers
  • Cosmos-1.0-Tokenizer-DV8x16x16 (8x temporal compression, 16x16 spatial compression, 720 short spatial size, 49 frames context)

The previous release (v0.1) of Cosmos Tokenizer included the following tokenizers:

• Continuous Tokenizers
  • Continuous Image (CI) Tokenizer
    • Cosmos-0.1-Tokenizer-CI8x8 (8x8 spatial compression)
    • Cosmos-0.1-Tokenizer-CI16x16 (16x16 spatial compression)
  • Continuous Video (CV) Tokenizer
    • Cosmos-0.1-Tokenizer-CV4x8x8 (4x temporal compression, 8x8 spatial compression)
    • Cosmos-0.1-Tokenizer-CV8x8x8 (8x temporal compression, 8x8 spatial compression)
    • Cosmos-0.1-Tokenizer-CV8x16x16 (8x temporal compression, 16x16 spatial compression)
• Discrete Tokenizers
  • Discrete Image (DI) Tokenizer
    • Cosmos-0.1-Tokenizer-DI8x8 (8x8 spatial compression)
    • Cosmos-0.1-Tokenizer-DI16x16 (16x16 spatial compression)
  • Discrete Video (DV) Tokenizer
    • Cosmos-0.1-Tokenizer-DV4x8x8 (4x temporal compression, 8x8 spatial compression)
    • Cosmos-1.0-Tokenizer-DV8x8x8 (8x temporal compression, 8x8 spatial compression)
    • Cosmos-1.0-Tokenizer-DV8x16x16 (8x temporal compression, 16x16 spatial compression)

License/Terms of Use:

NVIDIA Open Model License

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 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 (FSQ) as the latent space quantizer.

Input/Output Specifications

Encoder

• Input

  • Type: Images or Videos
  • Format: RGB (Red, Green, Blue)
  • Properties:
    • 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

  • Type: Tokens
  • Properties:
    • Continuous-valued feature vectors with a dimensionality of 16

Decoder

• Input

  • Type: Tokens
  • Properties:
    • Continuous-valued feature vectors with a dimensionality of 16

• Output

  • Type Images or Videos (matching input type)
  • Format: RGB (Red, Green, Blue)
  • Properties:
    • Resolution: Same as input resolution
    • Video Length: Same as input video length

Software Integration (Required For NVIDIA Models Only):

Runtime Engine(s):

• Cosmos-Tokenizer
• Cosmos-Inference

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 (PyTorch)
• 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.

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

• Install dependencies

  pip3 install -r requirements.txt
  apt-get install -y ffmpeg
  

• Preferably, you could build a docker image using our provided Dockerfile.

  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

• Under the checkpoint directory pretrained_ckpts/<model-name>, we provide the encoder, decoder and the full autoencoder JIT models.

  ├── pretrained_ckpts/  
  │   ├── Cosmos-Tokenizer-1.0-CV8x8x8/
  │   │   ├── 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.

import torch
from cosmos_tokenizer.video_lib import CausalVideoTokenizer

model_name = "Cosmos-Tokenizer-1.0-CV8x8x8"
input_tensor = torch.randn(1, 3, 9, 512, 512).to('cuda').to(torch.bfloat16)  # [B, C, T, H, W]
encoder = CausalVideoTokenizer(checkpoint_enc=f'pretrained_ckpts/{model_name}/encoder.jit')
(latent,) = encoder.encode(input_tensor)
torch.testing.assert_close(latent.shape, (1, 16, 2, 64, 64))

# The input tensor can be reconstructed by the decoder as:
decoder = CausalVideoTokenizer(checkpoint_dec=f'pretrained_ckpts/{model_name}/decoder.jit')
reconstructed_tensor = decoder.decode(latent)
torch.testing.assert_close(reconstructed_tensor.shape, input_tensor.shape)

The indices will have the shape (1, 2, 64, 64) and contain integral values in the range [1..64K], where the first of the three integral maps represents the first frame. The codes will contain the pre-quantization continuous latent with shape (1, 6, 2, 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.

Evaluation

TokenizationPerformance Comparison

We have evaluated the additional Cosmos Tokenizer models on DAVIS video benchmark dataset.

• We compare with the state-of-the-art discrete video tokenizer, OmniTokenizer.
• Evaluation metrics:
  • Peak Signal-to-Noise Ratio (PSNR)
  • Structural Similarity (SSIM)
  • Reconstruction Fréchet Video Distance (rFVD)

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.

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.