DOCSEARCHbyAlgolia
Submit Docs

TiTok PyTorch

repository·main·Indexed 41 days ago

https://github.com/lucidrains/titok-pytorch

A PyTorch implementation of a tokenizer that represents images using a small, fixed number of latent tokens instead of a dense spatial grid. Version 0.0.5 features the TiTokTokenizer class for efficient image compression and reconstruction via a transformer-based encoder-decoder architecture with vector quantization.

Tokens
1.1K
Snippets
5
Records
6
Agent score
2%

What's inside TiTok

  1. How TiTok tokenization works

    main

    Unlike standard Vision Transformers that tokenize every spatial patch, TiTok tokenizes a fixed set of num_latent_tokens.

    1. Patching: The image is divided into patches and projected to dim.
    2. Positioning: Absolute position embeddings are added to the patch tokens.
    3. Latent Concatenation: A set of learnable latents is concatenated with the patch tokens.
    4. Encoding: An Encoder processes the combined sequence.
    5. Quantization: The sequence is split, and only the latents portion is passed through the Vector Quantizer (vq). This produces the discrete codebook indices.
    6. Decoding: The quantized latents are passed through a Decoder and converted back into image patches to reconstruct the original image.

  2. Reconstruct images from codebook IDs using codebook_ids_to_images()

    main

    To reconstruct an image from discrete tokens, use the codebook_ids_to_images(token_ids) method. It retrieves the corresponding vectors from the codebook and passes them through the decoder.

    Input: token_ids tensor of shape (batch, num_latent_tokens) containing integer indices from the codebook.

    # token_ids shape: (batch, 32)
reconstructed_images = tokenizer.codebook_ids_to_images(token_ids)
# reconstructed_images shape: (batch, 3, 256, 256)

  3. Tokenize images using TiTokTokenizer.tokenize()

    main

    Use the .tokenize(images) method to convert an image tensor into discrete codebook indices. This method is a wrapper around .forward() with return_codebook_ids = True and is decorated with @torch.no_grad() to ensure no gradients are computed during inference.

    Input: images tensor of shape (batch, channels, height, width) where height and width must match image_size.

    # images shape: (batch, 3, 256, 256)
indices = tokenizer.tokenize(images)
# indices shape: (batch, num_latent_tokens)

  4. Initialize the TiTokTokenizer

    main

    The TiTokTokenizer class is a PyTorch module used for image tokenization and reconstruction. It uses a transformer-based encoder-decoder architecture with vector quantization (VQ) applied to a fixed number of latent tokens rather than spatial patches.

    Key parameters include:

    • dim: The embedding dimension.
    • image_size: The resolution of input images (default: 256).
    • patch_size: The size of the image patches (default: 32).
    • channels: Number of image channels (default: 3).
    • num_latent_tokens: The number of latent tokens to be quantized (default: 32).
    • codebook_size: The size of the vector quantization codebook (default: 8192).
    • enc_kwargs, dec_kwargs, vq_kwargs: Dictionaries passed to the underlying Encoder and VectorQuantize modules.
    from titok_pytorch import TiTokTokenizer
import torch

tokenizer = TiTokTokenizer(
    dim = 512,
    image_size = 256,
    patch_size = 32,
    num_latent_tokens = 32,
    codebook_size = 8192
)

  6. Use TiTokTokenizer.forward() for training

    main

    The .forward() method is designed for training the model. It computes the reconstruction loss and optionally returns the reconstructed images.

    Arguments:

    • images: Input image tensor.
    • return_codebook_ids: If True, returns only the quantized indices.
    • return_recon_images: If True, returns a tuple of (recon_loss, recon_images). If False, returns only recon_loss.

    Returns:

    • If return_recon_images=False: recon_loss (MSE loss between original and reconstructed images).
    • If return_recon_images=True: (recon_loss, recon_images).
    # For training (returns loss)
loss = tokenizer(images)

# For training (returns loss and images)
loss, recon_images = tokenizer(images, return_recon_images = True)

# To get indices only
indices = tokenizer(images, return_codebook_ids = True)