tokenizer method

README.md

@@ -12,10 +12,14 @@ from titok_pytorch.titok import TiTokTokenizer
 
 images = torch.randn(2, 3, 256, 256)
 
-tokenizer = TiTokTokenizer(dim = 512)
+titok = TiTokTokenizer(dim = 512)
 
-loss = tokenizer(images)
+loss = titok(images)
 loss.backward()
+
+# after much training
+
+codes = titok.tokenize(images)
 ```
 
 ## Citations

titok_pytorch/titok.py

@@ -98,9 +98,14 @@ def __init__(
             Rearrange('b h w (c p1 p2) -> b c (h p1) (w p2)', p1 = patch_size, p2 = patch_size)
         )
 
+    @torch.no_grad()
+    def tokenize(self, images):
+        return self.forward(images, return_codebook_ids = True)
+
     def forward(
         self,
-        images
+        images,
+        return_codebook_ids = False
     ):
         batch = images.shape[0]
         orig_images = images
@@ -136,6 +141,11 @@ def forward(
 
         quantized, indices, _ = self.vq(latents)
 
+        # whether to early return
+
+        if return_codebook_ids:
+            return indices
+
         # append mask tokens
 
         mask_tokens = repeat(self.mask_tokens, 'n d -> b n d', b = batch)