expose a few more vit hparams

README.md

@@ -19,9 +19,10 @@ from titok_pytorch import TiTokTokenizer
 images = torch.randn(2, 3, 256, 256)
 
 titok = TiTokTokenizer(
-    dim = 512,
+    dim = 1024,
+    patch_size = 32,
     num_latent_tokens = 32,   # they claim only 32 tokens needed
-    codebook_size = 8192      # codebook size 8192
+    codebook_size = 4096      # codebook size 4096
 )
 
 loss = titok(images)
@@ -42,8 +43,6 @@ assert recon_images.shape == images.shape
 ## Todo
 
 - [ ] add multi-resolution patches
-- [ ] add lfq
-- [ ] support video
 
 ## Citations
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titok-pytorch"
-version = "0.0.4"
+version = "0.0.5"
 description = "TiTok - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

titok_pytorch/titok.py

@@ -44,7 +44,11 @@ def __init__(
         channels = 3,
         num_latent_tokens = 32,
         enc_depth = 6,
+        enc_heads = 8,
+        enc_dim_head = 64,
         dec_depth = 6,
+        dec_heads = 8,
+        dec_dim_head = 64,
         codebook_size = 8192,
         enc_kwargs: dict = dict(),
         dec_kwargs: dict = dict(),
@@ -84,6 +88,8 @@ def __init__(
         self.encoder = Encoder(
             dim = dim,
             depth = enc_depth,
+            heads = enc_heads,
+            attn_dim_head = enc_dim_head,
             **enc_kwargs
         )
 
@@ -97,6 +103,8 @@ def __init__(
         self.decoder = Encoder(
             dim = dim,
             depth = dec_depth,
+            heads = dec_heads,
+            attn_dim_head = dec_dim_head,
             **dec_kwargs
         )