TiTok - Pytorch
Project description
TiTok - Pytorch (wip)
Implementation of TiTok, proposed by Bytedance in An Image is Worth 32 Tokens for Reconstruction and Generation
Install
$ pip install titok-pytorch
Usage
import torch
from titok_pytorch import TiTokTokenizer
images = torch.randn(2, 3, 256, 256)
titok = TiTokTokenizer(
dim = 1024,
patch_size = 32,
num_latent_tokens = 32, # they claim only 32 tokens needed
codebook_size = 4096 # codebook size 4096
)
loss = titok(images)
loss.backward()
# after much training
# extract codes for gpt, maskgit, whatever
codes = titok.tokenize(images) # (2, 32)
# reconstructing images from codes
recon_images = titok.codebook_ids_to_images(codes)
assert recon_images.shape == images.shape
Todo
- add multi-resolution patches
Citations
@article{yu2024an,
author = {Qihang Yu and Mark Weber and Xueqing Deng and Xiaohui Shen and Daniel Cremers and Liang-Chieh Chen},
title = {An Image is Worth 32 Tokens for Reconstruction and Generation},
journal = {arxiv: 2406.07550},
year = {2024}
}
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
titok_pytorch-0.0.5.tar.gz
(1.3 MB
view hashes)
Built Distribution
Close
Hashes for titok_pytorch-0.0.5-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | bf7e2c69c848c8386dc578ac2063191b2ffe757ac250f2b623013bc66d78f791 |
|
MD5 | adce9264bd54d91ca3aabf5fb70e816d |
|
BLAKE2b-256 | f7597098442f96023cd13ed9c9101148e069b5e3d1bfa8bbb86bac27278ae406 |