https://github.com/lucidrains/titok-pytorch
Implementation of TiTok, proposed by Bytedance in "An Image is Worth 32 Tokens for Reconstruction and Generation"
https://github.com/lucidrains/titok-pytorch
artificial-intelligence deep-learning image-generation vector-quantization
Last synced: over 1 year ago
JSON representation
Implementation of TiTok, proposed by Bytedance in "An Image is Worth 32 Tokens for Reconstruction and Generation"
- Host: GitHub
- URL: https://github.com/lucidrains/titok-pytorch
- Owner: lucidrains
- License: mit
- Created: 2024-06-17T18:59:25.000Z (about 2 years ago)
- Default Branch: main
- Last Pushed: 2024-06-20T13:41:52.000Z (about 2 years ago)
- Last Synced: 2025-03-31T16:14:23.252Z (over 1 year ago)
- Topics: artificial-intelligence, deep-learning, image-generation, vector-quantization
- Language: Python
- Homepage:
- Size: 1.25 MB
- Stars: 170
- Watchers: 9
- Forks: 5
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
## TiTok - Pytorch (wip)
Implementation of TiTok, proposed by Bytedance in An Image is Worth 32 Tokens for Reconstruction and Generation
## Install
```bash
$ pip install titok-pytorch
```
## Usage
```python
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
```bibtex
@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}
}
```