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 = 512,
num_latent_tokens = 32, # they claim only 32 tokens needed
codebook_size = 8192 # codebook size 8192
)
loss = titok(images)
loss.backward()
# after much training
# extract codes for gpt, maskgit, whatever
codes = titok.tokenize(images)
# reconstructing images from codes
recon_images = titok.codebook_ids_to_images(codes)
assert recon_images.shape == images.shape
Todo
- add multi-resolution patches
- add lfq
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.4.tar.gz
(1.3 MB
view hashes)
Built Distribution
Close
Hashes for titok_pytorch-0.0.4-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 41a1d12d53fb59846cb2b812adc2adce41113668d02f5711619df94c30487e08 |
|
MD5 | cf1ef667dbe866e3801c01bd0e3879d2 |
|
BLAKE2b-256 | 163b712b4f89de5d047aebdf4137c8860414849fed1e75f05fca56749fa58a69 |